diff --git a/Lib/cProfile.py b/Lib/cProfile.py --- a/Lib/cProfile.py +++ b/Lib/cProfile.py @@ -7,6 +7,7 @@ __all__ = ["run", "runctx", "Profile"] import _lsprof +import contextlib # ____________________________________________________________ # Simple interface @@ -56,6 +57,51 @@ result = prof.print_stats(sort) return result +def runcall(func, *args, filename=None, sort=-1): + """Run func(*args) under profiler, optionally saving results in + filename. + """ + prof = Profile() + result = None + try: + try: + prof.runcall(func, *args) + except SystemExit: + pass + finally: + if filename is not None: + prof.dump_stats(filename) + else: + result = prof.print_stats(sort) + return result + +@contextlib.contextmanager +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 + ... + >>> + """ + prof = Profile() + prof.enable() + try: + yield + finally: + prof.disable() + if filename is not None: + prof.dump_stats(filename) + else: + result = prof.print_stats(sort) + + # ____________________________________________________________ class Profile(_lsprof.Profiler): @@ -71,6 +117,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() @@ -129,20 +181,14 @@ return self.runctx(cmd, dict, dict) def runctx(self, cmd, globals, locals): - self.enable() - try: + with self: exec(cmd, globals, locals) - finally: - self.disable() return self # This method is more useful to profile a single function call. def runcall(self, func, *args, **kw): - self.enable() - try: + with self: return func(*args, **kw) - finally: - self.disable() # ____________________________________________________________ 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,44 @@ 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')