bpo-29235: Make cProfile.Profile a context manager (GH-6808)

4 files changed, 52 insertions, 0 deletions

diff --git a/Doc/library/profile.rst b/Doc/library/profile.rst
index a6dc56f..5dc0b2f 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 f166a1c..c804504 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 1430d22..2fd67ee 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 0000000..4618afc
--- /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.