Merge pull request #339 from TTsangSC/module-fix

FIX: `kernprof -m`: presence of the executed module as `sys.modules['__main__']`

Author: Erotemic

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
 ~~~~~

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):

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)

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: