Repository: Lightning-AI/pytorch-lightning
File: src/lightning/pytorch/profilers/profiler.py (Python)

Code before:

        Example::

            with self.profile('load training data'):
                # load training data code

        The profiler will start once you've entered the context and will automatically
        stop once you exit the code block.

        """
        try:
            self.start(action_name)
            yield action_name
        finally:
            self.stop(action_name)

    def _rank_zero_info(self, *args: Any, **kwargs: Any) -> None:
        if self._local_rank in (None, 0):
            log.info(*args, **kwargs)

    def _prepare_filename(
        self,
        action_name: Optional[str] = None,
        extension: str = ".txt",
        split_token: str = "-",  # noqa: S107
        sanitize: bool = True,
    ) -> str:
        args = []
        if self._stage is not None:
            args.append(self._stage)
        if self.filename:
            args.append(self.filename)
        if self._local_rank is not None:
            args.append(str(self._local_rank))
        if action_name is not None:
            args.append(action_name)
        base = split_token.join(args)
        if sanitize:
            # Replace a set of path-unsafe characters across platforms with '_'
            base = re.sub(r"[\\/:*?\"<>|\n\r\t]", "_", base)
            base = re.sub(r"_+", "_", base)
            base = base.strip()
            if not base:
                base = "profile"
        return base + extension

    def _prepare_streams(self) -> None:
        if self._write_stream is not None:
            return
        if self.filename and self.dirpath:
            filepath = os.path.join(self.dirpath, self._prepare_filename())

Review comment (security):
This looks like an internal flag, and it’s unclear how a user is meant to turn it off.

We should either sanitize unconditionally or treat this as a proper configuration option: pass it through the constructor and store it as an instance attribute so it can be used consistently later.

Code after:

        Example::

            with self.profile('load training data'):
                # load training data code

        The profiler will start once you've entered the context and will automatically
        stop once you exit the code block.

        """
        try:
            self.start(action_name)
            yield action_name
        finally:
            self.stop(action_name)

    def _rank_zero_info(self, *args: Any, **kwargs: Any) -> None:
        if self._local_rank in (None, 0):
            log.info(*args, **kwargs)

    def _prepare_filename(
        self,
        action_name: Optional[str] = None,
        extension: str = ".txt",
    ) -> str:
        args = []
        if self._stage is not None:
            args.append(self._stage)
        if self.filename:
            args.append(self.filename)
        if self._local_rank is not None:
            args.append(str(self._local_rank))
        if action_name is not None:
            args.append(action_name)
        base = "-".join(args)
        # Replace a set of path-unsafe characters across platforms with '_'
        base = re.sub(r"[\\/:*?\"<>|\n\r\t]", "_", base)
        base = re.sub(r"_+", "_", base)
        base = base.strip()
        if not base:
            base = "profile"
        return base + extension

    def _prepare_streams(self) -> None:
        if self._write_stream is not None:
            return
        if self.filename and self.dirpath:
            filepath = os.path.join(self.dirpath, self._prepare_filename())
            fs = get_filesystem(filepath)
            fs.mkdirs(self.dirpath, exist_ok=True)
            file = fs.open(filepath, "a")