From 0eb4aa1c95706a40b9fa83eeda38ec784b77b8c5 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Sun, 11 May 2025 02:28:03 +0200 Subject: [PATCH 1/7] added failing test --- tests/test_kernprof.py | 48 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/tests/test_kernprof.py b/tests/test_kernprof.py index cfb7a553..202a821f 100644 --- a/tests/test_kernprof.py +++ b/tests/test_kernprof.py @@ -1,3 +1,4 @@ +import contextlib import os import re import shlex @@ -8,7 +9,7 @@ import pytest import ubelt as ub -from kernprof import ContextualProfile +from kernprof import main, ContextualProfile def f(x): @@ -123,6 +124,51 @@ def main(): assert ('Function: main' in proc.stdout) == profiled_main +@pytest.mark.parametrize('error', [True, False]) +def test_kernprof_sys_restoration(capsys, error): + """ + Test that `kernprof.main()` properly restores `sys.path` on the way + out. + + Notes + ----- + The test is run in-process. + """ + with contextlib.ExitStack() as stack: + enter = stack.enter_context + tmpdir = enter(tempfile.TemporaryDirectory()) + assert tmpdir not in sys.path + temp_dpath = ub.Path(tmpdir) + (temp_dpath / 'mymod.py').write_text(ub.codeblock( + f''' + import sys + + + def main(): + # Mess up `sys.path` + sys.path.append({tmpdir!r}) + # Output + print(1) + # Optionally raise an error + if {error!r}: + raise Exception + + + if __name__ == '__main__': + main() + ''')) + enter(ub.ChDir(tmpdir)) + if error: + ctx = pytest.raises(BaseException) + else: + ctx = contextlib.nullcontext() + with ctx: + main(['-l', '-m', 'mymod']) + out, _ = capsys.readouterr() + assert out.startswith('1') + assert tmpdir not in sys.path + + class TestKernprof(unittest.TestCase): def test_enable_disable(self): From dba24005a125d46a1a8cfb52935b38566965bc53 Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Sun, 11 May 2025 02:30:21 +0200 Subject: [PATCH 2/7] Fixed `@contextlib.contextmanager` line_profiler/autoprofile/autoprofile.py::run() kernprof.py::main() tests/test_line_profiler.py::check_timings() Fixed bug in `@contextlib.contextmanager`s where the cleanup code is not called if an error occurred inside the `with` block --- kernprof.py | 6 ++++-- line_profiler/autoprofile/autoprofile.py | 8 +++++--- tests/test_line_profiler.py | 9 ++++++--- 3 files changed, 15 insertions(+), 8 deletions(-) diff --git a/kernprof.py b/kernprof.py index b8e558cf..0c1e89c7 100755 --- a/kernprof.py +++ b/kernprof.py @@ -249,8 +249,10 @@ def _restore_list(lst): [1, 2, 3] """ old = lst.copy() - yield - lst[:] = old + try: + yield + finally: + lst[:] = old def pre_parse_single_arg_directive(args, flag, sep='--'): diff --git a/line_profiler/autoprofile/autoprofile.py b/line_profiler/autoprofile/autoprofile.py index d5894318..2f0e2a33 100644 --- a/line_profiler/autoprofile/autoprofile.py +++ b/line_profiler/autoprofile/autoprofile.py @@ -100,9 +100,11 @@ def run(script_file, ns, prof_mod, profile_imports=False, as_module=False): @contextlib.contextmanager def restore_dict(d, target=None): copy = d.copy() - yield target - d.clear() - d.update(copy) + try: + yield target + finally: + d.clear() + d.update(copy) if as_module: Profiler = AstTreeModuleProfiler diff --git a/tests/test_line_profiler.py b/tests/test_line_profiler.py index b2a12028..9f2fb6a1 100644 --- a/tests/test_line_profiler.py +++ b/tests/test_line_profiler.py @@ -50,9 +50,12 @@ def check_timings(prof): timings = prof.get_stats().timings assert not any(timings.values()), ('Expected no timing entries, ' f'got {timings!r}') - yield prof - timings = prof.get_stats().timings - assert any(timings.values()), f'Expected timing entries, got {timings!r}' + try: + yield prof + finally: + timings = prof.get_stats().timings + assert (any(timings.values()), + f'Expected timing entries, got {timings!r}') def test_init(): From a102db22baa2a7f13d2b72fa14ffafd785ca0eff Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Sun, 11 May 2025 02:35:24 +0200 Subject: [PATCH 3/7] Added CHANGELOG entry --- CHANGELOG.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 62c53494..b9d96b1c 100644 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -10,6 +10,7 @@ Changes * 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+). +* FIX: Fixed ``@contextlib.contextmanager`` bug where the cleanup code (e.g. restoration of ``sys`` attributes) is not run if exceptions occurred inside the context 4.2.0 ~~~~~ From d6bca74e954b965307bc5fc9fee7c36450afc72c Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Sun, 11 May 2025 17:46:59 +0200 Subject: [PATCH 4/7] Refactoring: class-based context managers kernprof.py::_restore_list() line_profiler/autoprofile/autoprofile.py::run()..restore_dict tests/test_line_profiler.py::check_timings Refactored to use classes instead of `@contextlib.contextmanager` tests/test_kernprof.py::test_kernprof_sys_restoration() Extended to also check for the restoration of `sys.modules` tests/test_line_profiler.py::check_timings Fixed on-exit assertion, which was always true regardless of the timings because it was written `assert (cond, msg)` instead of `assert cond, (msg)` --- kernprof.py | 26 ++++++++++++++------- line_profiler/autoprofile/autoprofile.py | 21 ++++++++++------- tests/test_kernprof.py | 14 +++++++++--- tests/test_line_profiler.py | 29 +++++++++++++++--------- 4 files changed, 60 insertions(+), 30 deletions(-) diff --git a/kernprof.py b/kernprof.py index 0c1e89c7..a434691f 100755 --- a/kernprof.py +++ b/kernprof.py @@ -80,7 +80,6 @@ def main(): --prof-imports If specified, modules specified to `--prof-mod` will also autoprofile modules that they import. Only works with line_profiler -l, --line-by-line """ import builtins -import contextlib import functools import os import sys @@ -224,8 +223,7 @@ def _python_command(): return sys.executable -@contextlib.contextmanager -def _restore_list(lst): +class _restore_list: """ Restore a list like `sys.path` after running code which potentially modifies it. @@ -248,11 +246,23 @@ def _restore_list(lst): >>> l [1, 2, 3] """ - old = lst.copy() - try: - yield - finally: - lst[:] = old + def __init__(self, lst): + self.lst = lst + self.contexts = [] + + def __enter__(self): + self.contexts.append(self.lst.copy()) + + def __exit__(self, *_, **__): + self.lst[:] = self.contexts.pop() + + def __call__(self, func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + with self: + return func(*args, **kwargs) + + return wrapper def pre_parse_single_arg_directive(args, flag, sep='--'): diff --git a/line_profiler/autoprofile/autoprofile.py b/line_profiler/autoprofile/autoprofile.py index 2f0e2a33..04a1d6ef 100644 --- a/line_profiler/autoprofile/autoprofile.py +++ b/line_profiler/autoprofile/autoprofile.py @@ -97,14 +97,19 @@ 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 """ - @contextlib.contextmanager - def restore_dict(d, target=None): - copy = d.copy() - try: - yield target - finally: - d.clear() - d.update(copy) + class restore_dict: + def __init__(self, d, target=None): + self.d = d + self.target = target + self.contexts = [] + + def __enter__(self): + self.contexts.append(self.d.copy()) + return self.target + + def __exit__(self, *_, **__): + self.d.clear() + self.d.update(self.contexts.pop()) if as_module: Profiler = AstTreeModuleProfiler diff --git a/tests/test_kernprof.py b/tests/test_kernprof.py index 202a821f..c56dd215 100644 --- a/tests/test_kernprof.py +++ b/tests/test_kernprof.py @@ -125,9 +125,15 @@ def main(): @pytest.mark.parametrize('error', [True, False]) -def test_kernprof_sys_restoration(capsys, error): +@pytest.mark.parametrize( + 'args', + ['', '-pmymod'], # Normal execution / auto-profile +) +def test_kernprof_sys_restoration(capsys, error, args): """ - Test that `kernprof.main()` properly restores `sys.path` on the way + Test that `kernprof.main()` and + `line_profiler.autoprofile.autoprofile.run()` (resp.) properly + restores `sys.path` (resp. `sys.modules['__main__']`) on the way out. Notes @@ -162,11 +168,13 @@ def main(): ctx = pytest.raises(BaseException) else: ctx = contextlib.nullcontext() + old_main = sys.modules.get('__main__') with ctx: - main(['-l', '-m', 'mymod']) + main(['-l', *shlex.split(args), '-m', 'mymod']) out, _ = capsys.readouterr() assert out.startswith('1') assert tmpdir not in sys.path + assert sys.modules.get('__main__') is old_main class TestKernprof(unittest.TestCase): diff --git a/tests/test_line_profiler.py b/tests/test_line_profiler.py index 9f2fb6a1..85860b2a 100644 --- a/tests/test_line_profiler.py +++ b/tests/test_line_profiler.py @@ -41,21 +41,28 @@ def strip(s): return textwrap.dedent(s).strip('\n') -@contextlib.contextmanager -def check_timings(prof): +class check_timings: """ Verify that the profiler starts without timing data and ends with some. """ - timings = prof.get_stats().timings - assert not any(timings.values()), ('Expected no timing entries, ' - f'got {timings!r}') - try: - yield prof - finally: - timings = prof.get_stats().timings - assert (any(timings.values()), - f'Expected timing entries, got {timings!r}') + def __init__(self, prof): + self.prof = prof + + def __enter__(self): + timings = self.timings + assert not any(timings.values()), ( + f'Expected no timing entries, got {timings!r}') + return self.prof + + def __exit__(self, *_, **__): + timings = self.timings + assert any(timings.values()), ( + f'Expected timing entries, got {timings!r}') + + @property + def timings(self): + return self.prof.get_stats().timings def test_init(): From 336c7715914b72b5542dd6119fa6824af777eff0 Mon Sep 17 00:00:00 2001 From: Terence Tsang Date: Mon, 12 May 2025 18:20:21 +0200 Subject: [PATCH 5/7] Update kernprof.py Co-authored-by: Jon Crall --- kernprof.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/kernprof.py b/kernprof.py index a434691f..4fdc95a1 100755 --- a/kernprof.py +++ b/kernprof.py @@ -248,13 +248,14 @@ class _restore_list: """ def __init__(self, lst): self.lst = lst - self.contexts = [] + self.old = None def __enter__(self): - self.contexts.append(self.lst.copy()) + assert self.old is None + self.old = self.lst.copy() def __exit__(self, *_, **__): - self.lst[:] = self.contexts.pop() + self.lst[:] = self.old def __call__(self, func): @functools.wraps(func) From b8dab906b995db4b4d717c6c22b48f4f8727373d Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Mon, 12 May 2025 18:23:41 +0200 Subject: [PATCH 6/7] Fixed `kernprof.py::_restore_list()` --- kernprof.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/kernprof.py b/kernprof.py index 4fdc95a1..52867331 100755 --- a/kernprof.py +++ b/kernprof.py @@ -255,7 +255,7 @@ def __enter__(self): self.old = self.lst.copy() def __exit__(self, *_, **__): - self.lst[:] = self.old + self.old, self.lst[:] = None, self.old def __call__(self, func): @functools.wraps(func) From 7988790096ab3b64a3397614a05c7b4646f7e01b Mon Sep 17 00:00:00 2001 From: "Terence S.-C. Tsang" Date: Mon, 12 May 2025 18:29:16 +0200 Subject: [PATCH 7/7] Removed redundant handling for reentry --- line_profiler/autoprofile/autoprofile.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/line_profiler/autoprofile/autoprofile.py b/line_profiler/autoprofile/autoprofile.py index 04a1d6ef..5985a84b 100644 --- a/line_profiler/autoprofile/autoprofile.py +++ b/line_profiler/autoprofile/autoprofile.py @@ -101,15 +101,17 @@ class restore_dict: def __init__(self, d, target=None): self.d = d self.target = target - self.contexts = [] + self.copy = None def __enter__(self): - self.contexts.append(self.d.copy()) + assert self.copy is None + self.copy = self.d.copy() return self.target def __exit__(self, *_, **__): self.d.clear() - self.d.update(self.contexts.pop()) + self.d.update(self.copy) + self.copy = None if as_module: Profiler = AstTreeModuleProfiler