diff --git a/Lib/cProfile.py b/Lib/cProfile.py --- a/Lib/cProfile.py +++ b/Lib/cProfile.py @@ -8,6 +8,7 @@ import _lsprof import profile as _pyprofile +import contextlib # ____________________________________________________________ # Simple interface @@ -19,8 +20,17 @@ return _pyprofile._Utils(Profile).runctx(statement, globals, locals, filename, sort) +def runcall(func, *args, filename=None, sort=-1): + return _pyprofile._Utils(Profile).runcall(func, *args, filename=filename, + sort=sort) + +def runblock(filename=None, sort=-1): + return _pyprofile._Utils(Profile).runblock(filename, sort) + run.__doc__ = _pyprofile.run.__doc__ runctx.__doc__ = _pyprofile.runctx.__doc__ +runcall.__doc__ = _pyprofile.runcall.__doc__ +runblock.__doc__ = _pyprofile.runblock.__doc__ # ____________________________________________________________ @@ -37,6 +47,12 @@ # Most of the functionality is in the base class. # This subclass only adds convenient and backward-compatible methods. + def __enter__(self): + self.enable() + + def __exit__(self, *args): + self.disable() + def print_stats(self, sort=-1): import pstats pstats.Stats(self).strip_dirs().sort_stats(sort).print_stats() @@ -110,6 +126,14 @@ finally: self.disable() + @contextlib.contextmanager + def runblock(self): + self.enable() + try: + yield self + finally: + self.disable() + # ____________________________________________________________ def label(code): diff --git a/Lib/profile.py b/Lib/profile.py --- a/Lib/profile.py +++ b/Lib/profile.py @@ -28,6 +28,7 @@ import os import time import marshal +import contextlib from optparse import OptionParser __all__ = ["run", "runctx", "Profile"] @@ -67,6 +68,29 @@ finally: self._show(prof, filename, sort) + def runcall(self, func, *args, filename=None, sort=-1): + prof = self.profiler() + try: + prof.runcall(func, *args) + except SystemExit: + pass + finally: + self._show(prof, filename, sort) + + @contextlib.contextmanager + def runblock(self, filename=None, sort=-1): + prof = self.profiler() + try: + with prof.runblock(): + yield + except SystemExit: + pass + finally: + if filename is not None: + prof.dump_stats(filename) + else: + prof.print_stats(sort) + def _show(self, prof, filename, sort): if filename is not None: prof.dump_stats(filename) @@ -100,6 +124,28 @@ """ return _Utils(Profile).runctx(statement, globals, locals, filename, sort) +def runcall(func, *args, filename=None, sort=-1): + """Run func(*args) under profiler, optionally saving results in + filename. + """ + return _Utils(Profile).runcall(func, *args, filename=filename, sort=sort) + +def runblock(filename=None, sort=-1): + """Function that runs a block of code under profile, and can be + used as a context manager or a decorator. + Example: + + >>> with runblock(): + ... pass + ... + >>> @runblock() + ... def foo(): + ... pass + ... + >>> + """ + return _Utils(Profile).runblock(filename, sort) + class Profile: """Profiler class. @@ -435,6 +481,14 @@ finally: sys.setprofile(None) + @contextlib.contextmanager + def runblock(self): + self.set_cmd('') + sys.setprofile(self.dispatcher) + try: + yield self + finally: + sys.setprofile(None) #****************************************************************** # The following calculates the overhead for using a profiler. The diff --git a/Lib/test/test_profile.py b/Lib/test/test_profile.py --- a/Lib/test/test_profile.py +++ b/Lib/test/test_profile.py @@ -93,6 +93,43 @@ filename=TESTFN) self.assertTrue(os.path.exists(TESTFN)) + def test_runcall(self): + flag = [] + with silent(): + self.profilermodule.runcall(flag.append, 1) + self.assertEqual(flag, [1]) + self.profilermodule.runcall(flag.append, 2, filename=TESTFN) + self.assertEqual(flag, [1, 2]) + self.assertTrue(os.path.exists(TESTFN)) + + def test_runblock_ctx_manager(self): + flag = [] + with silent(): + with self.profilermodule.runblock(): + flag.append(1) + self.assertEqual(flag, [1]) + with self.profilermodule.runblock(filename=TESTFN): + flag.append(2) + self.assertEqual(flag, [1, 2]) + self.assertTrue(os.path.exists(TESTFN)) + + def test_runblock_decorator(self): + flag = [] + + @self.profilermodule.runblock() + def foo(): + flag.append(1) + with silent(): + foo() + self.assertEqual(flag, [1]) + + @self.profilermodule.runblock(filename=TESTFN) + def foo(): + flag.append(1) + with silent(): + foo() + self.assertEqual(flag, [1, 1]) + self.assertTrue(os.path.exists(TESTFN)) def regenerate_expected_output(filename, cls): filename = filename.rstrip('co')