From: Scott Sanderson <ssanderson@quantopian.com>
Date: Fri, 1 Jun 2018 20:36:23 +0000 (-0400)
Subject: bpo-29235: Make cProfile.Profile a context manager (GH-6808)
X-Git-Tag: v3.8.0a1~1693
X-Git-Url: https://granicus.if.org/sourcecode?a=commitdiff_plain;h=2e01b75884892d5aabdaab658fbd17f7a7ccebaa;p=python

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

diff --git a/Doc/library/profile.rst b/Doc/library/profile.rst
index a6dc56f43c..5dc0b2f940 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 f166a1c437..c8045043cb 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 1430d22504..2fd67ee756 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 0000000000..4618afc13a
--- /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.