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

Author: Scott Sanderson

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.

Lib/cProfile.py

@@ -110,6 +110,13 @@ def runcall(self, func, *args, **kw):
         finally:
             self.disable()
 
+    def __enter__(self):
+        self.enable()
+        return self
+
+    def __exit__(self, *exc_info):
+        self.disable()
+
 # ____________________________________________________________
 
 def label(code):

Lib/test/test_cprofile.py

@@ -49,6 +49,33 @@ def test_module_path_option(self):
         # 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)

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.