From 5068a32dd19ce3f71bb8a326ca805eb8ab958951 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Sat, 18 Jul 2026 08:49:06 +0200 Subject: [PATCH 1/3] Test for #436 tests/test_kernprof.py chdir_temp New helper context manager for creating a new tempdir and `cd`-ing into it, factored out from common pattern in tests test_*() Updated to use the above test_kernprof_m_import_resolution() test_kernprof_eager_preimport_bad_module() Replaced hard-coded ':' in the `${PYTHONPATH}` value with `os.pathsep` test_ppe_pickling() New test which uses `concurrent.futures.ProcessPoolExecutor` in code run by `kernprof.main()`; the case where PPE hasn't been imported prior to `kernprof.main()` being called twice currently fails --- tests/test_kernprof.py | 166 ++++++++++++++++++++++++++++++++--------- 1 file changed, 129 insertions(+), 37 deletions(-) diff --git a/tests/test_kernprof.py b/tests/test_kernprof.py index 19a53d91..a36018cc 100644 --- a/tests/test_kernprof.py +++ b/tests/test_kernprof.py @@ -26,6 +26,41 @@ def g(x): yield y + 20 +class chdir_temp: + """ + Helper context for ``chdir``-ing to a tempdir. + + Example: + >>> from pathlib import Path + + >>> get_cwd = lambda: Path.cwd().resolve() + >>> cwd = get_cwd() + >>> ctx = chdir_temp() + >>> with ctx as new_cwd: + ... assert new_cwd.samefile(get_cwd()) + ... assert not new_cwd.samefile(cwd) + ... with ctx as newer_cwd: # Reentrance + ... assert newer_cwd.samefile(get_cwd()) + ... assert not newer_cwd.samefile(new_cwd) + ... assert not newer_cwd.samefile(cwd) + ... assert new_cwd.samefile(get_cwd()) + ... + >>> assert cwd.samefile(get_cwd()) + """ + def __init__(self) -> None: + self._stacks: list[contextlib.ExitStack] = [] + + def __enter__(self) -> ub.Path: + stack = contextlib.ExitStack() + self._stacks.append(stack) + tmpdir = stack.enter_context(tempfile.TemporaryDirectory()) + stack.enter_context(ub.ChDir(tmpdir)) + return ub.Path(tmpdir) + + def __exit__(self, *_, **__) -> None: + self._stacks.pop().close() + + @pytest.mark.parametrize( 'use_kernprof_exec, args, expected_output, expect_error', [ @@ -68,8 +103,7 @@ def test_kernprof_m_parsing( an argument and cuts off everything after it, passing that along to the module to be executed. """ - with tempfile.TemporaryDirectory() as tmpdir: - temp_dpath = ub.Path(tmpdir) + with chdir_temp() as temp_dpath: mod = (temp_dpath / 'mymod.py').resolve() mod.write_text( ub.codeblock( @@ -86,7 +120,7 @@ def test_kernprof_m_parsing( cmd = ['kernprof'] else: cmd = [sys.executable, '-m', 'kernprof'] - proc = ub.cmd(cmd + args, cwd=temp_dpath, verbose=2) + proc = ub.cmd(cmd + args, verbose=2) if expect_error: assert proc.returncode return @@ -116,8 +150,7 @@ def test_kernprof_m_sys_modules(flags, profiled_main): Test that `kernprof -m` is amenable to modules relying on the global `sys` state (e.g. those using `@enum.global_enum`). """ - with tempfile.TemporaryDirectory() as tmpdir: - temp_dpath = ub.Path(tmpdir) + with chdir_temp() as temp_dpath: (temp_dpath / 'mymod.py').write_text( ub.codeblock( """ @@ -152,7 +185,7 @@ def main(): '-m', 'mymod', ] - proc = ub.cmd(cmd, cwd=temp_dpath, verbose=2) + proc = ub.cmd(cmd, verbose=2) proc.check_returncode() assert proc.stdout.startswith('3\n') assert ('Function: main' in proc.stdout) == profiled_main @@ -188,14 +221,13 @@ def main(): line for line in code.splitlines() if '@profile' not in line ) cmd += ['-p', 'my_namesapce_pkg.mysubmod'] - with tempfile.TemporaryDirectory() as tmpdir: - temp_dpath = ub.Path(tmpdir) + with chdir_temp() as temp_dpath: namespace_mod_path = temp_dpath / 'my_namesapce_pkg' / 'mysubmod.py' namespace_mod_path.parent.mkdir() namespace_mod_path.write_text(code) - python_path = tmpdir + python_path = str(temp_dpath) if 'PYTHONPATH' in os.environ: - python_path += ':' + os.environ['PYTHONPATH'] + python_path += os.pathsep + os.environ['PYTHONPATH'] env = { **os.environ, # Toggle use of static analysis @@ -204,7 +236,7 @@ def main(): 'PYTHONPATH': python_path, } cmd += ['-m', 'my_namesapce_pkg.mysubmod'] - proc = ub.cmd(cmd, cwd=temp_dpath, verbose=2, env=env) + proc = ub.cmd(cmd, verbose=2, env=env) if static: assert proc.returncode assert proc.stderr.startswith('Could not find module') @@ -230,11 +262,9 @@ def test_kernprof_sys_restoration(capsys, error, args): ----- The test is run in-process. """ - with contextlib.ExitStack() as stack: - enter = stack.enter_context - tmpdir = enter(tempfile.TemporaryDirectory()) + with chdir_temp() as temp_dpath: + tmpdir = str(temp_dpath) assert tmpdir not in sys.path - temp_dpath = ub.Path(tmpdir) (temp_dpath / 'mymod.py').write_text( ub.codeblock( f""" @@ -256,7 +286,6 @@ def main(): """ ) ) - enter(ub.ChDir(tmpdir)) if error: ctx = pytest.raises(BaseException) else: @@ -345,10 +374,7 @@ def test_kernprof_verbosity(flags, expected_stdout, expected_stderr): """ Test the various verbosity levels of `kernprof`. """ - with contextlib.ExitStack() as stack: - enter = stack.enter_context - tmpdir = enter(tempfile.TemporaryDirectory()) - temp_dpath = ub.Path(tmpdir) + with chdir_temp() as temp_dpath: (temp_dpath / 'script.py').write_text( ub.codeblock( """ @@ -365,7 +391,6 @@ def main(): """ ) ) - enter(ub.ChDir(tmpdir)) proc = ub.cmd( [ 'kernprof', @@ -408,17 +433,13 @@ def test_kernprof_eager_preimport_bad_module(): in an auto-generated pre-import module. """ bad_module = """raise Exception('Boo')""" - with contextlib.ExitStack() as stack: - enter = stack.enter_context - tmpdir = enter(tempfile.TemporaryDirectory()) - temp_dpath = ub.Path(tmpdir) + with chdir_temp() as temp_dpath: (temp_dpath / 'my_bad_module.py').write_text(bad_module) - enter(ub.ChDir(tmpdir)) python_path = os.environ.get('PYTHONPATH', '') if python_path: - python_path = f'{python_path}:{tmpdir}' + python_path = f'{python_path}{os.pathsep}{temp_dpath}' else: - python_path = tmpdir + python_path = str(temp_dpath) proc = ub.cmd( [ 'kernprof', @@ -454,9 +475,7 @@ def test_kernprof_bad_temp_script(stdin): in a temporary script supplied via `kernprof -c` or `kernprof -`. """ bad_script = """1 / 0""" - with contextlib.ExitStack() as stack: - enter = stack.enter_context - enter(ub.ChDir(enter(tempfile.TemporaryDirectory()))) + with chdir_temp(): if stdin: proc = subprocess.run( ['kernprof', '-'], @@ -490,9 +509,7 @@ def test_bad_prof_mod_target(debug): """ Test the handling of bad paths in `--prof-mod` targets. """ - with contextlib.ExitStack() as stack: - enter = stack.enter_context - enter(ub.ChDir(enter(tempfile.TemporaryDirectory()))) + with chdir_temp(): proc = ub.cmd( [ 'kernprof', @@ -519,9 +536,7 @@ def test_call_with_diagnostics(module, builtin): Test the output of call signatures in debug messages. """ to_run = ['-m', 'calendar'] if module else ['-c', 'print("Output: foo")'] - with contextlib.ExitStack() as stack: - enter = stack.enter_context - enter(ub.ChDir(enter(tempfile.TemporaryDirectory()))) + with chdir_temp(): cmd = ['kernprof'] if builtin: cmd += ['-b'] @@ -542,6 +557,83 @@ def test_call_with_diagnostics(module, builtin): assert bool(has_execfile_call) == (not module) +@pytest.mark.parametrize('preimport_ppe', [True, False]) +@pytest.mark.parametrize('n', [1, 2]) +def test_ppe_pickling(n: int, preimport_ppe: bool): + """ + Test that code using + :py:class:`concurrent.futures.ProcessPoolExecutor` doesn't cause + :py:func:`kernprof.main` to fail when called more than once. + + See also: + Issue #436 + """ + if preimport_ppe: + preimport = 'from concurrent.futures import ProcessPoolExecutor' + else: # Something irrelevant + preimport = 'from sys import modules # noqa' + + inner_script = ub.codeblock(""" + from __future__ import annotations + + from multiprocessing import get_context + from concurrent.futures import ProcessPoolExecutor + + + def my_sum(x: list[int]) -> int: + result = 0 + for n in x: + result += n + return result + + + def main() -> None: + with ProcessPoolExecutor(mp_context=get_context('spawn')) as ex: + print(list(ex.map(my_sum, [[1, 2], [3, 4], [5, 6]]))) + + + if __name__ == '__main__': + main() + """) + outer_script = ub.codeblock(f""" + import os + {preimport} + from tempfile import TemporaryDirectory + + from kernprof import main as kp_main + + + def main() -> None: + for _ in range({n}): + with TemporaryDirectory() as tmpdir: + kp_main([ + '-lzv', + f'--outfile={{os.path.join(tmpdir, "out.lprof")}}', + '--prof-mod=inner_script.py', + 'inner_script.py', + ]) + + + if __name__ == '__main__': + main() + """) + with chdir_temp() as temp_dpath: + (temp_dpath / 'inner_script.py').write_text(inner_script) + (temp_dpath / 'outer_script.py').write_text(outer_script) + proc = subprocess.run( + [sys.executable, 'outer_script.py'], + capture_output=True, text=True, + ) + for stream in 'stdout', 'stderr': + content = getattr(proc, stream) + if content is None: + continue + print(content, file=getattr(sys, stream)) + proc.check_returncode() + assert proc.stdout.startswith('[3, 7, 11]') + assert 'def main' in proc.stdout + + class TestKernprof(unittest.TestCase): def test_enable_disable(self): profile = ContextualProfile() From f6b18de1b8f0ba4e9d154e715cbf2ae847207678 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Sat, 18 Jul 2026 05:04:21 +0200 Subject: [PATCH 2/3] Fix #436 kernprof.py _restore Migrated to `line_profiler.line_profiler_utils.restore` main() Instead of reverting the entire module namespace of `line_profiler.diagnostics`, now only reverting `line_profiler.diagnostics.log` (which is replaced in `_parse_arguments()`) _main_profile() Instead of reverting the entire `sys.modules`, now only reverting `sys.modules['__main__']` line_profiler/autoprofile/autoprofile.py::run() Replaced internal context manager with a `restore` instance line_profiler/line_profiler_utils.py::restore - Migrated from `kernprof.py::_restore` - Added type annotations - Added optional argument `keys` (resp. `attrs`) to `.mapping()` (resp. `.instance_dict()`) so that users can selectively revert certain key-value pairs (resp. attributes) - Now reentrant - Note: the attributes `.setter`, `.getter`, and `.obj` are removed --- CHANGELOG.rst | 2 + kernprof.py | 119 +------------- line_profiler/autoprofile/autoprofile.py | 19 +-- line_profiler/line_profiler_utils.py | 191 ++++++++++++++++++++++- 4 files changed, 196 insertions(+), 135 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 30156f63..b91c64da 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -12,6 +12,8 @@ Changes is available (#427) * FIX: Bytecodes of profiled functions now always labeled to prevent confusion with non-profiled "twins" (#425) +* FIX: Stop reverting ``sys.modules`` after calling ``kernprof.main()`` + to avoid edge-case issues with e.g. pickling (#437) 5.0.2 diff --git a/kernprof.py b/kernprof.py index 8a7c4d6a..6650c410 100755 --- a/kernprof.py +++ b/kernprof.py @@ -228,6 +228,9 @@ def main(): positive_float, short_string_path, ) +from line_profiler.line_profiler_utils import ( + restore as _restore, # Compatibility +) from line_profiler.profiler_mixin import ByCountProfilerMixin from line_profiler._logger import Logger from line_profiler import _diagnostics as diagnostics @@ -404,118 +407,6 @@ def find(path): return list(results) -class _restore: - """ - Restore a collection like :py:data:`sys.path` after running code - which potentially modifies it. - """ - - def __init__(self, obj, getter, setter): - self.obj = obj - self.setter = setter - self.getter = getter - self.old = None - - def __enter__(self): - assert self.old is None - self.old = self.getter(self.obj) - - def __exit__(self, *_, **__): - self.setter(self.obj, self.old) - self.old = None - - def __call__(self, func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - with self: - return func(*args, **kwargs) - - return wrapper - - @classmethod - def sequence(cls, seq): - """ - Example - ------- - >>> l = [1, 2, 3] - >>> - >>> with _restore.sequence(l): - ... print(l) - ... l.append(4) - ... print(l) - ... l[:] = 5, 6 - ... print(l) - ... - [1, 2, 3] - [1, 2, 3, 4] - [5, 6] - >>> l - [1, 2, 3] - """ - - def set_list(orig, copy): - orig[:] = copy - - return cls(seq, methodcaller('copy'), set_list) - - @classmethod - def mapping(cls, mpg): - """ - Example - ------- - >>> d = {1: 2} - >>> - >>> with _restore.mapping(d): - ... print(d) - ... d[2] = 3 - ... print(d) - ... d.clear() - ... d.update({1: 4, 3: 5}) - ... print(d) - ... - {1: 2} - {1: 2, 2: 3} - {1: 4, 3: 5} - >>> d - {1: 2} - """ - - def set_mapping(orig, copy): - orig.clear() - orig.update(copy) - - return cls(mpg, methodcaller('copy'), set_mapping) - - @classmethod - def instance_dict(cls, obj): - """ - Example - ------- - >>> class Obj: - ... def __init__(self, x, y): - ... self.x, self.y = x, y - ... - ... def __repr__(self): - ... return 'Obj({0.x!r}, {0.y!r})'.format(self) - ... - >>> - >>> obj = Obj(1, 2) - >>> - >>> with _restore.instance_dict(obj): - ... print(obj) - ... obj.x, obj.y, obj.z = 4, 5, 6 - ... print(obj, obj.z) - ... - Obj(1, 2) - Obj(4, 5) 6 - >>> obj - Obj(1, 2) - >>> hasattr(obj, 'z') - False - """ - return cls.mapping(vars(obj)) - - def pre_parse_single_arg_directive(args, flag, sep='--'): """ Pre-parse high-priority single-argument directives like @@ -924,7 +815,7 @@ def _parse_arguments( @_restore.sequence(sys.argv) @_restore.sequence(sys.path) -@_restore.instance_dict(diagnostics) +@_restore.instance_dict(diagnostics, ['log']) def main(args=None, *, exit_on_error=True): """ Runs the command line interface @@ -1372,7 +1263,7 @@ def _main_profile(options, module=False, exit_on_error=True): runner, target = 'execfile', script_file assert runner in module_ns - with _restore.mapping(sys.modules): + with _restore.mapping(sys.modules, ['__main__']): sys.modules['__main__'] = module_obj if options.builtin: call(module_ns[runner], target, module_ns) diff --git a/line_profiler/autoprofile/autoprofile.py b/line_profiler/autoprofile/autoprofile.py index e287d35c..4471a778 100644 --- a/line_profiler/autoprofile/autoprofile.py +++ b/line_profiler/autoprofile/autoprofile.py @@ -52,6 +52,7 @@ def main(): from collections.abc import MutableMapping from typing import Any, cast, Dict, Mapping from typing import ContextManager +from ..line_profiler_utils import restore from .ast_tree_profiler import AstTreeProfiler from .run_module import AstTreeModuleProfiler from .line_profiler_utils import add_imported_function_or_module @@ -107,22 +108,6 @@ def run( as_module (bool): Whether we're running script_file as a module """ - - class restore_dict: - def __init__(self, d: MutableMapping[str, Any]): - self.d = d - self.copy: Mapping[str, Any] | None = None - - def __enter__(self): - assert self.copy is None - self.copy = dict(self.d) - - def __exit__(self, *_, **__): - self.d.clear() - if self.copy is not None: - self.d.update(self.copy) - self.copy = None - Profiler: type[AstTreeModuleProfiler] | type[AstTreeProfiler] if as_module: @@ -148,7 +133,7 @@ def __exit__(self, *_, **__): _extend_line_profiler_for_profiling_imports(ns[PROFILER_LOCALS_NAME]) code_obj = compile(tree_profiled, script_file, 'exec') - with restore_dict(sys.modules): + with restore.mapping(sys.modules, ['__main__']): # Always set the module object to `sys.modules['__main__']` and # then restore it via the context manager, so that the executed # code is run as `__main__` diff --git a/line_profiler/line_profiler_utils.py b/line_profiler/line_profiler_utils.py index 887cdd55..706df475 100644 --- a/line_profiler/line_profiler_utils.py +++ b/line_profiler/line_profiler_utils.py @@ -5,10 +5,21 @@ from __future__ import annotations import enum -import typing +from collections.abc import ( + Callable, Collection, Mapping, MutableMapping, MutableSequence, Sequence, +) +from functools import wraps +from operator import methodcaller +from typing import TYPE_CHECKING, Any, Generic, TypeVar, final +from typing_extensions import Self, ParamSpec -if typing.TYPE_CHECKING: - from typing_extensions import Self + +T = TypeVar('T') +K = TypeVar('K') +V = TypeVar('V') +T1 = TypeVar('T1') +T2 = TypeVar('T2') +PS = ParamSpec('PS') class _StrEnumBase(str, enum.Enum): @@ -49,7 +60,7 @@ def __str__(self) -> str: try: from enum import StrEnum as _StrEnum except ImportError: - if not typing.TYPE_CHECKING: # Don't confuse the typechecker + if not TYPE_CHECKING: # Don't confuse the typechecker _StrEnum = _StrEnumBase @@ -89,3 +100,175 @@ def _missing_(cls, value: object) -> Self | None: for name, instance in cls.__members__.items() } return members.get(value.casefold()) + + +@final +class restore(Generic[T1, T2]): + """ + Context manager for restoring a collection like :py:data:`sys.path` + after running code which potentially modifies it. + + Notes: + - Mainly to be used via the class-method instantiators. + + - Also permits being used as a decorator. + + - Instances are reentrant (see e.g. the doctest of + :py:meth:`.sequence`). + """ + def __init__( + self, + obj: T1, + getter: Callable[[T1], T2], + setter: Callable[[T1, T2], Any], + ) -> None: + self.obj = obj + self._setter = setter + self._getter = getter + self._stack: list[T2] = [] + + def __enter__(self) -> None: + self._stack.append(self._getter(self.obj)) + + def __exit__(self, *_, **__) -> None: + self._setter(self.obj, self._stack.pop()) + + def __call__(self, func: Callable[PS, T]) -> Callable[PS, T]: + @wraps(func) + def wrapper(*args: PS.args, **kwargs: PS.kwargs) -> T: + with self: + return func(*args, **kwargs) + + return wrapper + + @classmethod + def sequence( + cls: type[restore], seq: MutableSequence[T], + ) -> restore[MutableSequence[T], MutableSequence[T]]: + """ + Example: + >>> l = [1, 2, 3] + >>> + >>> restore_list = restore.sequence(l) + >>> with restore_list: + ... print(l) + ... l.append(4) + ... print(l) + ... with restore_list: # Reentrance + ... l[:] = 5, 6 + ... print(l) + ... print(l) + ... + [1, 2, 3] + [1, 2, 3, 4] + [5, 6] + [1, 2, 3, 4] + >>> l + [1, 2, 3] + """ + + def set_list(orig: MutableSequence[T], copy: Sequence[T]) -> None: + orig[:] = copy + + return cls(seq, methodcaller('copy'), set_list) + + @classmethod + def mapping( + cls: type[restore], + mpg: MutableMapping[K, V], + keys: Collection[K] | None = None, + ) -> restore[MutableMapping[K, V], Mapping[K, V | Any]]: + """ + Example + ------- + Whole-dict preservation: + + >>> d = {1: 2} + + >>> with restore.mapping(d): + ... print(d) + ... d[2] = 3 + ... print(d) + ... d.clear() + ... d.update({1: 4, 3: 5}) + ... print(d) + ... + {1: 2} + {1: 2, 2: 3} + {1: 4, 3: 5} + >>> d + {1: 2} + + Only preserving select key-value pairs: + + >>> d = {1: 2, 3: 4} + + >>> with restore.mapping(d, [1, 2]): + ... print(d) + ... d[1], d[2], d[3], d[4] = 5, 6, 7, 8 + ... print(d) + {1: 2, 3: 4} + {1: 5, 3: 7, 2: 6, 4: 8} + >>> d # [1] reverted to 2, [2] reverted to none + {1: 2, 3: 7, 4: 8} + """ + + def set_mapping( + orig: MutableMapping[K, V], copy: Mapping[K, V], + ) -> None: + orig.clear() + orig.update(copy) + + def get_subset(orig: Mapping[K, V]) -> Mapping[K, V | Sentinel]: + if TYPE_CHECKING: + assert keys is not None + return {k: orig.get(k, sentinel) for k in keys} + + def set_subset( + orig: MutableMapping[K, V], subset: Mapping[K, V | Sentinel], + ) -> None: + for key, value in subset.items(): + if value is sentinel: + orig.pop(key, None) + else: + orig[key] = value + + class Sentinel(enum.Enum): + sentinel = enum.auto() + + sentinel = Sentinel.sentinel + if keys is None: + return cls(mpg, methodcaller('copy'), set_mapping) + else: + return cls(mpg, get_subset, set_subset) + + @classmethod + def instance_dict( + cls: type[restore], obj: Any, attrs: Collection[str] | None = None, + ) -> restore[MutableMapping[str, Any], Mapping[str, Any]]: + """ + Example + ------- + >>> class Obj: + ... def __init__(self, x, y): + ... self.x, self.y = x, y + ... + ... def __repr__(self): + ... return 'Obj({0.x!r}, {0.y!r})'.format(self) + ... + >>> + >>> obj = Obj(1, 2) + >>> + >>> with restore.instance_dict(obj): + ... print(obj) + ... obj.x, obj.y, obj.z = 4, 5, 6 + ... print(obj, obj.z) + ... + Obj(1, 2) + Obj(4, 5) 6 + >>> obj + Obj(1, 2) + >>> hasattr(obj, 'z') + False + """ + return cls.mapping(vars(obj), attrs) From 5845332b4cb54c5616cfb82c4c88d48aa33293c1 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Sat, 18 Jul 2026 17:22:08 +0200 Subject: [PATCH 3/3] Misc changes kernprof.py No longer importing `line_profiler.line_profiler_utils.restore` under the name `_restore` since it isn't fully compatible to the old `_restore` anyway tests/test_kernprof.py::test_ppe_pickling() Now explicitly passing `max_workers` to PPE to avoid spinning up too many processes --- kernprof.py | 12 +++++------- tests/test_kernprof.py | 4 +++- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/kernprof.py b/kernprof.py index 6650c410..590fd48c 100755 --- a/kernprof.py +++ b/kernprof.py @@ -228,9 +228,7 @@ def main(): positive_float, short_string_path, ) -from line_profiler.line_profiler_utils import ( - restore as _restore, # Compatibility -) +from line_profiler.line_profiler_utils import restore from line_profiler.profiler_mixin import ByCountProfilerMixin from line_profiler._logger import Logger from line_profiler import _diagnostics as diagnostics @@ -813,9 +811,9 @@ def _parse_arguments( return options, tempfile_source_and_content -@_restore.sequence(sys.argv) -@_restore.sequence(sys.path) -@_restore.instance_dict(diagnostics, ['log']) +@restore.sequence(sys.argv) +@restore.sequence(sys.path) +@restore.instance_dict(diagnostics, ['log']) def main(args=None, *, exit_on_error=True): """ Runs the command line interface @@ -1263,7 +1261,7 @@ def _main_profile(options, module=False, exit_on_error=True): runner, target = 'execfile', script_file assert runner in module_ns - with _restore.mapping(sys.modules, ['__main__']): + with restore.mapping(sys.modules, ['__main__']): sys.modules['__main__'] = module_obj if options.builtin: call(module_ns[runner], target, module_ns) diff --git a/tests/test_kernprof.py b/tests/test_kernprof.py index a36018cc..ce679c0f 100644 --- a/tests/test_kernprof.py +++ b/tests/test_kernprof.py @@ -588,7 +588,9 @@ def my_sum(x: list[int]) -> int: def main() -> None: - with ProcessPoolExecutor(mp_context=get_context('spawn')) as ex: + with ProcessPoolExecutor( + max_workers=2, mp_context=get_context('spawn'), + ) as ex: print(list(ex.map(my_sum, [[1, 2], [3, 4], [5, 6]])))