diff --git a/Doc/library/profile.rst b/Doc/library/profile.rst index a6dc56f43cb..5dc0b2f9405 100644 --- a/Doc/library/profile.rst +++ b/Doc/library/profile.rst @@ -262,6 +262,16 @@ functions: ps.print_stats() print(s.getvalue()) + The :class:`Profile` class can also be used as a context manager (see + :ref:`typecontextmanager`):: + + import cProfile + + with cProfile.Profile() as pr: + # ... do something ... + + pr.print_stats() + .. method:: enable() Start collecting profiling data. diff --git a/Lib/cProfile.py b/Lib/cProfile.py index f166a1c4375..c8045043cbc 100755 --- a/Lib/cProfile.py +++ b/Lib/cProfile.py @@ -110,6 +110,13 @@ class Profile(_lsprof.Profiler): finally: self.disable() + def __enter__(self): + self.enable() + return self + + def __exit__(self, *exc_info): + self.disable() + # ____________________________________________________________ def label(code): diff --git a/Lib/test/test_cprofile.py b/Lib/test/test_cprofile.py index 1430d225048..2fd67ee7568 100644 --- a/Lib/test/test_cprofile.py +++ b/Lib/test/test_cprofile.py @@ -49,6 +49,33 @@ class CProfileTest(ProfileTest): # Test successful run assert_python_ok('-m', 'cProfile', '-m', 'timeit', '-n', '1') + def test_profile_enable_disable(self): + prof = self.profilerclass() + # Make sure we clean ourselves up if the test fails for some reason. + self.addCleanup(prof.disable) + + prof.enable() + self.assertIs(sys.getprofile(), prof) + + prof.disable() + self.assertIs(sys.getprofile(), None) + + def test_profile_as_context_manager(self): + prof = self.profilerclass() + # Make sure we clean ourselves up if the test fails for some reason. + self.addCleanup(prof.disable) + + with prof as __enter__return_value: + # profile.__enter__ should return itself. + self.assertIs(prof, __enter__return_value) + + # profile should be set as the global profiler inside the + # with-block + self.assertIs(sys.getprofile(), prof) + + # profile shouldn't be set once we leave the with-block. + self.assertIs(sys.getprofile(), None) + def test_main(): run_unittest(CProfileTest) diff --git a/Misc/NEWS.d/next/Library/2018-05-14-15-01-55.bpo-29235.47Fzwt.rst b/Misc/NEWS.d/next/Library/2018-05-14-15-01-55.bpo-29235.47Fzwt.rst new file mode 100644 index 00000000000..4618afc13ae --- /dev/null +++ b/Misc/NEWS.d/next/Library/2018-05-14-15-01-55.bpo-29235.47Fzwt.rst @@ -0,0 +1,8 @@ +The :class:`cProfile.Profile` class can now be used as a context manager. +You can profile a block of code by running:: + + import cProfile + with cProfile.Profile() as profiler: + # ... code to be profiled ... + +Patch by Scott Sanderson.