Shuffled code inside ~._child_process_profiling

line_profiler/_child_process_profiling/
    - The functions `pth_hook.py::_setup_in_child_process()` and
      `::_wrap_os_fork()` have been relocated to eponymous instance
      methods of `cache.py::LineProfilingCache`
    - The implementations of `pth_hook.py::{write,load}_pth_hook()` and
      `multiprocessing_patches.py::PickleHook.__setstate__()` are
      updated accordingly

Author: TTsangSC

line_profiler/_child_process_profiling/cache.py

@@ -4,14 +4,15 @@
 """
 from __future__ import annotations
 
+import atexit
 import dataclasses
 import os
 try:
     import _pickle as pickle
 except ImportError:
     import pickle  # type: ignore[assignment,no-redef]
 from collections.abc import Collection, Callable
-from functools import partial, cached_property
+from functools import partial, cached_property, wraps
 from operator import setitem
 from pathlib import Path
 from pickle import HIGHEST_PROTOCOL
@@ -20,6 +21,14 @@
 from typing_extensions import Self, ParamSpec
 
 from .. import _diagnostics as diagnostics
+from ..autoprofile.autoprofile import (
+    # Note: we need this to equip the profiler with the
+    # `.add_imported_function_or_module()` pseudo-method
+    # (see `kernprof.py::_write_preimports()`), which is required for
+    # the preimports to work
+    _extend_line_profiler_for_profiling_imports as upgrade_profiler,
+)
+from ..curated_profiling import CuratedProfilerContext
 from ..line_profiler import LineProfiler, LineStats
 # Note: this should have been defined here in this file, but we moved it
 # over to `~._child_process_hook` because that module contains the .pth
@@ -227,6 +236,116 @@ def _debug_output(self, msg: str) -> None:
         except OSError:  # Cache dir may have been rm-ed during cleanup
             pass
 
+    def _setup_in_child_process(
+        self,
+        wrap_os_fork: bool = False,
+        context: str = '',
+        prof: LineProfiler | None = None,
+    ) -> bool:
+        """
+        Set up shop in a forked/spawned child process so that
+        (line-)profiling can extend therein.
+
+        Args:
+            wrap_os_fork (bool):
+                Whether to wrap :py:func:`os.fork` which handles
+                profiling; already-forked child processes should set
+                this to false
+            context (str):
+                Optional context from which the function is called, to
+                be used in log messages
+            prof (LineProfiler | None):
+                Optional profiler instance to associate with the cache;
+                if not provided, an instance is created
+
+        Returns:
+            has_set_up (bool):
+                False the instance has already been set up prior to
+                calling this function, true otherwise
+        """
+        if not context:
+            context = '...'
+        self._debug_output(f'Setting up ({context})...')
+        if self.profiler is not None:  # Already set up
+            self._debug_output(f'Setup aborted ({context})')
+            return False
+
+        # Create a profiler instance and manage it with
+        # `CuratedProfilerContext`
+        if prof is None:
+            prof = LineProfiler()
+        self.profiler = prof
+        upgrade_profiler(prof)
+        ctx = CuratedProfilerContext(prof, insert_builtin=self.insert_builtin)
+        ctx.install()
+        self.add_cleanup(ctx.uninstall)
+        self._debug_output(f'Set up `.profiler` at {id(prof):#x}')
+
+        # Do the preimports at `cache.preimports_module` where
+        # appropriate
+        if self.preimports_module:
+            self._debug_output('Loading preimports...')
+            with open(self.preimports_module, mode='rb') as fobj:
+                code = compile(fobj.read(), self.preimports_module, 'exec')
+                exec(code, {})  # Use a fresh, empty namespace
+
+        # Occupy a tempfile slot in `.cache_dir` and set the profiler
+        # up to write thereto when the process terminates (with high
+        # priority)
+        prof_outfile = self.make_tempfile(
+            prefix='child-prof-output-{}-{}-{:#x}-'
+            .format(self.main_pid, os.getpid(), id(prof)),
+            suffix='.lprof',
+        )
+        self._add_cleanup(prof.dump_stats, -1, prof_outfile)
+
+        # Set up `os.fork()` wrapping if needed (i.e. in a spawned
+        # process)
+        if wrap_os_fork:
+            self._wrap_os_fork()
+
+        # Set `.cleanup()` as an atexit hook to handle everything when
+        # the child process is about to terminate
+        atexit.register(self.cleanup)
+
+        self._debug_output(f'Setup successful ({context})')
+        return True
+
+    def _wrap_os_fork(self) -> None:
+        """
+        Create a wrapper around :py:func:`os.fork` which handles
+        profiling.
+
+        Side effects:
+            - :py:func:`os.fork` (if available) replaced with the
+              wrapper
+            - :py:meth:`~.cleanup` callback registered undoing that
+        """
+        try:
+            fork = os.fork
+        except AttributeError:  # Can't fork on this platform
+            return
+
+        @wraps(fork)
+        def wrapper() -> int:
+            result = fork()
+            if result:
+                return result
+            # If we're here, we are in the fork
+            forked = self.copy()  # Ditch inherited cleanups
+            if forked._replace_loaded_instance():
+                forked._debug_output(
+                    'Superseded cached `.load()`-ed instance in forked process'
+                )
+            # Note: we can reuse the profiler instance in the fork, but
+            # it needs to go through setup so that the separate
+            # profiling results are dumped into another output file
+            forked._setup_in_child_process(False, 'fork', self.profiler)
+            return result
+
+        os.fork = wrapper
+        self.add_cleanup(setattr, os, 'fork', fork)
+
     def make_tempfile(self, **kwargs) -> Path:
         """
         Create a fresh tempfile under :py:attr:`~.cache_dir`. The other

line_profiler/_child_process_profiling/multiprocessing_patches.py

@@ -28,7 +28,6 @@
 
 from .. import _diagnostics as diagnostics
 from .cache import LineProfilingCache
-from .pth_hook import _setup_in_child_process
 from .runpy_patches import create_runpy_wrapper
 
 
@@ -80,7 +79,7 @@ def __setstate__(*_) -> None:
         # We're in a child process created by `multiprocessing`, so set
         # up shop here.
         lp_cache = LineProfilingCache.load()
-        _setup_in_child_process(lp_cache, False, 'multiprocessing')
+        lp_cache._setup_in_child_process(False, 'multiprocessing')
         # In a child process, we don't care about polluting the
         # `multiprocessing` namespace, so don't bother with cleanup
         if not getattr(multiprocessing, _PATCHED_MARKER, False):

line_profiler/_child_process_profiling/pth_hook.py

@@ -20,16 +20,11 @@
 from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
-    from collections.abc import Callable  # noqa: F401
     from pathlib import Path  # noqa: F401
-    from typing import Any  # noqa: F401
-    from ..line_profiler import LineProfiler  # noqa: F401
     from .cache import LineProfilingCache  # noqa: F401
 
 
-__all__ = (
-    'write_pth_hook', 'load_pth_hook', '_setup_in_child_process'
-)
+__all__ = ('write_pth_hook', 'load_pth_hook')
 
 INHERITED_PID_ENV_VARNAME = (
     'LINE_PROFILER_PROFILE_CHILD_PROCESSES_CACHE_PID'
@@ -83,7 +78,7 @@ def write_pth_hook(cache):  # type: (LineProfilingCache) -> Path
     finally:  # Not closing the handle causes issues on Windows
         os.close(handle)
 
-    _wrap_os_fork(cache)
+    cache._wrap_os_fork()
 
     return fpath
 
@@ -118,135 +113,11 @@ def load_pth_hook(ppid):  # type: (int) -> None
         return
     try:
         cache = LineProfilingCache.load()
-        _setup_in_child_process(cache, True, 'pth')
-        # _setup_in_child_process(LineProfilingCache.load())
+        cache._setup_in_child_process(True, 'pth')
     except Exception as e:
         if DEBUG:
             msg = f'{type(e)}: {e}'
             warnings.warn(msg)
             log.warning(msg)
     finally:
         load_pth_hook.called = True  # type: ignore[attr-defined]
-
-
-def _wrap_os_fork(cache):  # type: (LineProfilingCache) -> None
-    """
-    Create a wrapper around :py:func:`os.fork` which handles profiling.
-
-    Args:
-        cache (:py:class:`~.LineProfilingCache`):
-            Cache object
-
-    Side effects:
-        - :py:func:`os.fork` (if available) replaced with the wrapper
-        - Cleanup callback registered at ``cache`` undoing that
-    """
-    import os
-    from functools import wraps
-
-    try:
-        fork = os.fork
-    except AttributeError:  # Can't fork on this platform
-        return
-
-    @wraps(fork)
-    def wrapper():  # type: () -> int
-        result = fork()
-        if result:
-            return result
-        # If we're here, we are in the fork
-        forked = cache.copy()
-        if forked._replace_loaded_instance():
-            forked._debug_output(
-                'Superseded cached `.load()`-ed instance in forked process'
-            )
-        # Note: we can reuse the profiler instance in the fork, but it
-        # needs to go through setup so that the separate profiling
-        # results are dumped into another output file
-        _setup_in_child_process(forked, False, 'fork', cache.profiler)
-        return result
-
-    os.fork = wrapper
-    cache.add_cleanup(setattr, os, 'fork', fork)
-
-
-def _setup_in_child_process(cache, wrap_os_fork=False, context='', prof=None):
-    # type: (LineProfilingCache, bool, str, LineProfiler | None) -> bool
-    """
-    Set up shop in a forked/spawned child process so that
-    (line-)profiling can extend therein.
-
-    Args:
-        cache (LineProfilingCache):
-            Cache object
-        wrap_os_fork (bool):
-            Whether to wrap :py:func:`os.fork` which handles profiling;
-            already-forked child processes should set this to false
-        context (str):
-            Optional context from which the function is called, to be
-            used in log messages
-        prof (LineProfiler | None):
-            Optional profiler instance to associate with the cache;
-            if not provided, an instance is created
-
-    Returns:
-        has_set_up (bool):
-            False is ``cache`` has already been set up prior to calling
-            this function, true otherwise
-    """
-    if not context:
-        context = '...'
-    cache._debug_output(f'Setting up ({context})...')
-    if cache.profiler is not None:  # Already set up
-        cache._debug_output(f'Setup aborted ({context})')
-        return False
-
-    import os
-    from atexit import register
-    from ..autoprofile.autoprofile import (
-        # Note: we need this to equip the profiler with the
-        # `.add_imported_function_or_module()` pseudo-method
-        # (see `kernprof.py::_write_preimports()`), which is required
-        # for the preimports to work
-        _extend_line_profiler_for_profiling_imports as upgrade_profiler,
-    )
-    from ..curated_profiling import CuratedProfilerContext
-    from ..line_profiler import LineProfiler  # noqa: F811
-
-    # Create a profiler instance and manage it with
-    # `CuratedProfilerContext`
-    if prof is None:
-        prof = LineProfiler()
-    cache.profiler = prof
-    upgrade_profiler(prof)
-    ctx = CuratedProfilerContext(prof, insert_builtin=cache.insert_builtin)
-    ctx.install()
-    cache.add_cleanup(ctx.uninstall)
-    cache._debug_output(f'Set up `.profiler` at {id(prof):#x}')
-
-    # Do the preimports at `cache.preimports_module` where appropriate
-    if cache.preimports_module:
-        cache._debug_output('Loading preimports...')
-        with open(cache.preimports_module, mode='rb') as fobj:
-            code = compile(fobj.read(), cache.preimports_module, 'exec')
-            exec(code, {})  # Use a fresh, empty namespace
-
-    # Occupy a tempfile slot in `cache.cache_dir` and set the profiler
-    # up to write thereto when the process terminates
-    prof_outfile = cache.make_tempfile(
-        prefix='child-prof-output-{}-{}-{:#x}-'
-        .format(cache.main_pid, os.getpid(), id(prof)),
-        suffix='.lprof',
-    )
-    cache._add_cleanup(prof.dump_stats, -1, prof_outfile)
-
-    # Set up `os.fork()` wrapping if needed (i.e. in a spawned process)
-    if wrap_os_fork:
-        _wrap_os_fork(cache)
-
-    # Set `cache.cleanup()` as an atexit hook to handle everything when
-    # the child process is about to terminate
-    register(cache.cleanup)
-
-    cache._debug_output(f'Setup successful ({context})')
-    return True