diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 0e242a51..62c53494 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -9,6 +9,7 @@ Changes * FIX: Fixed auto-profiling of async function definitions #330 * ENH: Added CLI argument ``-m`` to ``kernprof`` for running a library module as a script; also made it possible for profiling targets to be supplied across multiple ``-p`` flags * FIX: Fixed explicit profiling of class methods; added handling for profiling static, bound, and partial methods, ``functools.partial`` objects, (cached) properties, and async generator functions +* FIX: Fixed namespace bug when running ``kernprof -m`` on certain modules (e.g. ``calendar`` on Python 3.12+). 4.2.0 ~~~~~ diff --git a/kernprof.py b/kernprof.py index ef096486..b8e558cf 100755 --- a/kernprof.py +++ b/kernprof.py @@ -470,7 +470,8 @@ def positive_float(value): try: try: execfile_ = execfile - rmod_ = run_module + rmod_ = functools.partial(run_module, + run_name='__main__', alter_sys=True) ns = locals() if options.prof_mod and options.line_by_line: from line_profiler.autoprofile import autoprofile @@ -487,13 +488,11 @@ def positive_float(value): profile_imports=options.prof_imports, as_module=module is not None) elif module and options.builtin: - run_module(options.script, ns, '__main__') + rmod_(options.script, ns) elif options.builtin: execfile(script_file, ns, ns) elif module: - prof.runctx(f'rmod_({options.script!r}, globals(), "__main__")', - ns, - ns) + prof.runctx(f'rmod_({options.script!r}, globals())', ns, ns) else: prof.runctx('execfile_(%r, globals())' % (script_file,), ns, ns) except (KeyboardInterrupt, SystemExit): diff --git a/line_profiler/autoprofile/autoprofile.py b/line_profiler/autoprofile/autoprofile.py index 7b738c86..d5894318 100644 --- a/line_profiler/autoprofile/autoprofile.py +++ b/line_profiler/autoprofile/autoprofile.py @@ -44,11 +44,16 @@ def main(): python -m kernprof -p demo.py -l demo.py python -m line_profiler -rmt demo.py.lprof """ - +import contextlib +import functools +import importlib.util +import operator +import sys import types from .ast_tree_profiler import AstTreeProfiler from .run_module import AstTreeModuleProfiler from .line_profiler_utils import add_imported_function_or_module +from .util_static import modpath_to_modname PROFILER_LOCALS_NAME = 'prof' @@ -92,10 +97,42 @@ def run(script_file, ns, prof_mod, profile_imports=False, as_module=False): as_module (bool): Whether we're running script_file as a module """ - Profiler = AstTreeModuleProfiler if as_module else AstTreeProfiler + @contextlib.contextmanager + def restore_dict(d, target=None): + copy = d.copy() + yield target + d.clear() + d.update(copy) + + if as_module: + Profiler = AstTreeModuleProfiler + module_name = modpath_to_modname(script_file) + if not module_name: + raise ModuleNotFoundError(f'script_file = {script_file!r}: ' + 'cannot find corresponding module') + + module_obj = types.ModuleType(module_name) + namespace = vars(module_obj) + namespace.update(ns) + + # Set the `__spec__` correctly + module_obj.__spec__ = importlib.util.find_spec(module_name) + + # Set the module object to `sys.modules` via a callback, and + # then restore it via the context manager + callback = functools.partial( + operator.setitem, sys.modules, '__main__', module_obj) + ctx = restore_dict(sys.modules, callback) + else: + Profiler = AstTreeProfiler + namespace = ns + ctx = contextlib.nullcontext(lambda: None) + profiler = Profiler(script_file, prof_mod, profile_imports) tree_profiled = profiler.profile() _extend_line_profiler_for_profiling_imports(ns[PROFILER_LOCALS_NAME]) code_obj = compile(tree_profiled, script_file, 'exec') - exec(code_obj, ns, ns) + with ctx as callback: + callback() + exec(code_obj, namespace, namespace) diff --git a/tests/test_kernprof.py b/tests/test_kernprof.py index 98a1e5cb..cfb7a553 100644 --- a/tests/test_kernprof.py +++ b/tests/test_kernprof.py @@ -1,3 +1,6 @@ +import os +import re +import shlex import sys import tempfile import unittest @@ -24,14 +27,14 @@ def g(x): 'use_kernprof_exec, args, expected_output, expect_error', [([False, ['-m'], '', True]), # `python -m kernprof` - (False, ['-m', 'mymod'], "['mymod']", False), + (False, ['-m', 'mymod'], "[__MYMOD__]", False), # `kernprof` - (True, ['-m', 'mymod'], "['mymod']", False), - (False, ['-m', 'mymod', '-p', 'bar'], "['mymod', '-p', 'bar']", False), + (True, ['-m', 'mymod'], "[__MYMOD__]", False), + (False, ['-m', 'mymod', '-p', 'bar'], "[__MYMOD__, '-p', 'bar']", False), # `-p bar` consumed by `kernprof`, `-p baz` are not (False, ['-p', 'bar', '-m', 'mymod', '-p', 'baz'], - "['mymod', '-p', 'baz']", + "[__MYMOD__, '-p', 'baz']", False), # Separator `--` broke off the remainder, so the requisite arg for # `-m` is not found and we error out @@ -49,26 +52,75 @@ def test_kernprof_m_parsing( an argument and cuts off everything after it, passing that along to the module to be executed. """ - temp_dpath = ub.Path(tempfile.mkdtemp()) - (temp_dpath / 'mymod.py').write_text(ub.codeblock( - ''' - import sys - - - if __name__ == '__main__': - print(sys.argv) - ''')) - if use_kernprof_exec: - cmd = ['kernprof'] - else: - cmd = [sys.executable, '-m', 'kernprof'] - proc = ub.cmd(cmd + args, cwd=temp_dpath, verbose=2) + with tempfile.TemporaryDirectory() as tmpdir: + temp_dpath = ub.Path(tmpdir) + mod = (temp_dpath / 'mymod.py').resolve() + mod.write_text(ub.codeblock( + ''' + import sys + + + if __name__ == '__main__': + print(sys.argv) + ''')) + if use_kernprof_exec: + cmd = ['kernprof'] + else: + cmd = [sys.executable, '-m', 'kernprof'] + proc = ub.cmd(cmd + args, cwd=temp_dpath, verbose=2) if expect_error: assert proc.returncode return else: proc.check_returncode() - assert proc.stdout.startswith(expected_output) + expected_output = re.escape(expected_output).replace( + '__MYMOD__', "'.*{}'".format(re.escape(os.path.sep + mod.name))) + assert re.match('^' + expected_output, proc.stdout) + + +@pytest.mark.skipif(sys.version_info[:2] < (3, 11), + reason='no `@enum.global_enum` in Python ' + f'{".".join(str(v) for v in sys.version_info[:3])}') +@pytest.mark.parametrize(('flags', 'profiled_main'), + [('-lv -p mymod', True), # w/autoprofile + ('-lv', True), # w/o autoprofile + ('-b', False)]) # w/o line profiling +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) + (temp_dpath / 'mymod.py').write_text(ub.codeblock( + ''' + import enum + import os + import sys + + + @enum.global_enum + class MyEnum(enum.Enum): + FOO = 1 + BAR = 2 + + + @profile + def main(): + x = FOO.value + BAR.value + print(x) + assert x == 3 + + + if __name__ == '__main__': + main() + ''')) + cmd = [sys.executable, '-m', 'kernprof', + *shlex.split(flags), '-m', 'mymod'] + proc = ub.cmd(cmd, cwd=temp_dpath, verbose=2) + proc.check_returncode() + assert proc.stdout.startswith('3\n') + assert ('Function: main' in proc.stdout) == profiled_main class TestKernprof(unittest.TestCase): @@ -130,6 +182,7 @@ def test_gen_decorator(self): next(i) self.assertEqual(profile.enable_count, 0) + if __name__ == '__main__': """ CommandLine: