Skip to content

FIX: Don't revert the entirety of sys.modules after running kernprof.main #437

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Erotemic merged 3 commits into from
Jul 18, 2026
Merged

FIX: Don't revert the entirety of sys.modules after running kernprof.main #437

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5068a32
Test for #436
TTsangSC Jul 18, 2026
f6b18de
Fix #436
TTsangSC Jul 18, 2026
5845332
Misc changes
TTsangSC Jul 18, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions CHANGELOG.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
121 changes: 5 additions & 116 deletions kernprof.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -228,6 +228,7 @@ def main():
positive_float,
short_string_path,
)
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
Expand Down Expand Up @@ -404,118 +405,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
Expand Down Expand Up @@ -922,9 +811,9 @@ def _parse_arguments(
return options, tempfile_source_and_content


@_restore.sequence(sys.argv)
@_restore.sequence(sys.path)
@_restore.instance_dict(diagnostics)
@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
Expand Down Expand Up @@ -1372,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):
with restore.mapping(sys.modules, ['__main__']):
sys.modules['__main__'] = module_obj
if options.builtin:
call(module_ns[runner], target, module_ns)
Expand Down
19 changes: 2 additions & 17 deletions line_profiler/autoprofile/autoprofile.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand All @@ -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__`
Expand Down
Loading
Loading