[Python-checkins] bpo-29235: Make cProfile.Profile a context manager (GH-6808)

[Brett Cannon]

https://github.com/python/cpython/commit/2e01b75884892d5aabdaab658fbd17f7a7ccebaa
commit: 2e01b75884892d5aabdaab658fbd17f7a7ccebaa
branch: master
author: Scott Sanderson <ssanderson at quantopian.com>
committer: Brett Cannon <brettcannon at users.noreply.github.com>
date: 2018-06-01T13:36:23-07:00
summary:

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

files:
A Misc/NEWS.d/next/Library/2018-05-14-15-01-55.bpo-29235.47Fzwt.rst
M Doc/library/profile.rst
M Lib/cProfile.py
M Lib/test/test_cprofile.py

diff --git a/Doc/library/profile.rst b/Doc/library/profile.rst
index a6dc56f43cbc..5dc0b2f94053 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 f166a1c43752..c8045043cbcf 100755
--- a/Lib/cProfile.py
+++ b/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):
diff --git a/Lib/test/test_cprofile.py b/Lib/test/test_cprofile.py
index 1430d2250485..2fd67ee75688 100644
--- a/Lib/test/test_cprofile.py
+++ b/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)
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 000000000000..4618afc13ae0
--- /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.