gh-132737: Support profiling modules that require __main___ (#132738)

This commit is contained in:
Aneesh Durg 2025-04-24 13:41:01 -05:00 committed by GitHub
parent e1c09fff05
commit c7a7aa9a57
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 29 additions and 3 deletions

View File

@ -6,6 +6,7 @@
import _lsprof
import importlib.machinery
import importlib.util
import io
import profile as _pyprofile
@ -173,13 +174,22 @@ def main():
code = compile(fp.read(), progname, 'exec')
spec = importlib.machinery.ModuleSpec(name='__main__', loader=None,
origin=progname)
globs = {
module = importlib.util.module_from_spec(spec)
# Set __main__ so that importing __main__ in the profiled code will
# return the same namespace that the code is executing under.
sys.modules['__main__'] = module
# Ensure that we're using the same __dict__ instance as the module
# for the global variables so that updates to globals are reflected
# in the module's namespace.
globs = module.__dict__
globs.update({
'__spec__': spec,
'__file__': spec.origin,
'__name__': spec.name,
'__package__': None,
'__cached__': None,
}
})
try:
runctx(code, globs, None, options.outfile, options.sort)
except BrokenPipeError as exc:

View File

@ -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
@ -154,6 +156,19 @@ def test_sort(self):
self.assertGreater(rc, 0)
self.assertIn(b"option -s: invalid choice: 'demo'", err)
def test_profile_script_importing_main(self):
"""Check that scripts that reference __main__ see their own namespace
when being profiled."""
with tempfile.NamedTemporaryFile("w+", delete_on_close=False) as f:
f.write(textwrap.dedent("""\
class Foo:
pass
import __main__
assert Foo == __main__.Foo
"""))
f.close()
assert_python_ok('-m', "cProfile", f.name)
def main():
if '-r' not in sys.argv:

View File

@ -0,0 +1 @@
Support profiling code that requires ``__main__``, such as :mod:`pickle`.