gh-132737: Support profiling modules that import __main___

Lib/cProfile.py

@@ -6,6 +6,7 @@
 
 import _lsprof
 import importlib.machinery
+import importlib.util
 import io
 import profile as _pyprofile
 
@@ -162,10 +163,8 @@ def main():
     if len(args) > 0:
         if options.module:
             code = "run_module(modname, run_name='__main__')"
-            globs = {
-                'run_module': runpy.run_module,
-                'modname': args[0]
-            }
+            globs = globals().copy()
+            globs.update({"run_module": runpy.run_module, "modname": args[0]})
         else:
             progname = args[0]
             sys.path.insert(0, os.path.dirname(progname))
@@ -179,10 +178,21 @@ def main():
                 '__name__': spec.name,
                 '__package__': None,
                 '__cached__': None,
+                '__builtins__': __builtins__,
             }
+        # cmd has to run in __main__ namespace (or imports from __main__ will
+        # break). Clear __main__ and replace with the globals provided.
+        import __main__
+        # Save a reference to the current __main__ namespace so that we can
+        # restore it after cmd completes.
+        original_main = __main__.__dict__.copy()
+        __main__.__dict__.update(globs)
+
         try:
-            runctx(code, globs, None, options.outfile, options.sort)
+            runctx(code, __main__.__dict__, None, options.outfile, options.sort)
         except BrokenPipeError as exc:
+            __main__.__dict__.clear()
+            __main__.__dict__.update(original_main)
             # Prevent "Exception ignored" during interpreter shutdown.
             sys.stdout = None
             sys.exit(exc.errno)

Lib/test/test_cprofile.py

@@ -5,8 +5,10 @@
 
 # rip off all interesting stuff from test_profile
 import cProfile
+import tempfile
+import textwrap
 from test.test_profile import ProfileTest, regenerate_expected_output
-from test.support.script_helper import assert_python_failure
+from test.support.script_helper import assert_python_failure, assert_python_ok
 from test import support
 
 
@@ -155,6 +157,22 @@ def test_sort(self):
         self.assertIn(b"option -s: invalid choice: 'demo'", err)
 
 
+class TestProfilingScript(unittest.TestCase):
+    def test_profile_script_importing_main(self):
+        """Check that scripts that reference __main__ see their own namespace
+        when being profiled."""
+        with tempfile.NamedTemporaryFile("w+") as f:
+            f.write(textwrap.dedent("""\
+                class Foo:
+                    pass
+
+                import __main__
+                assert Foo == __main__.Foo
+                """))
+            f.flush()
+            assert_python_ok('-m', "cProfile", f.name)
+
+
 def main():
     if '-r' not in sys.argv:
         unittest.main()

Misc/NEWS.d/next/Core_and_Builtins/2025-04-19-18-07-34.gh-issue-132737.9mW1il.rst

@@ -0,0 +1 @@
+Support profiling modules that import ``__main__``, such as modules that use pickle.