Issue #11647: allow contextmanager objects to be used as decorators as described in the docs. Initial patch by Ysj Ray.

6 files changed, 50 insertions, 10 deletions

diff --git a/Doc/library/contextlib.rst b/Doc/library/contextlib.rst
index a35ea569c1..e8dc17fed3 100644
--- a/Doc/library/contextlib.rst
+++ b/Doc/library/contextlib.rst
@@ -54,8 +54,12 @@ Functions provided:
    the exception has been handled, and execution will resume with the statement
    immediately following the :keyword:`with` statement.
 
-   contextmanager uses :class:`ContextDecorator` so the context managers it
-   creates can be used as decorators as well as in :keyword:`with` statements.
+   :func:`contextmanager` uses :class:`ContextDecorator` so the context managers
+   it creates can be used as decorators as well as in :keyword:`with` statements.
+   When used as a decorator, a new generator instance is implicitly created on
+   each function call (this allows the otherwise "one-shot" context managers
+   created by :func:`contextmanager` to meet the requirement that context
+   managers support multiple invocations in order to be used as decorators).
 
    .. versionchanged:: 3.2
       Use of :class:`ContextDecorator`.
@@ -155,6 +159,12 @@ Functions provided:
           def __exit__(self, *exc):
               return False
 
+   .. note::
+      As the decorated function must be able to be called multiple times, the
+      underlying context manager must support use in multiple :keyword:`with`
+      statements. If this is not the case, then the original construct with the
+      explicit :keyword:`with` statement inside the function should be used.
+
    .. versionadded:: 3.2
 
 
diff --git a/Lib/contextlib.py b/Lib/contextlib.py
index 4633cff7a4..e90fc4181a 100644
--- a/Lib/contextlib.py
+++ b/Lib/contextlib.py
@@ -9,10 +9,23 @@ __all__ = ["contextmanager", "closing", "ContextDecorator"]
 
 class ContextDecorator(object):
     "A base class or mixin that enables context managers to work as decorators."
+
+    def _recreate_cm(self):
+        """Return a recreated instance of self.
+        
+        Allows otherwise one-shot context managers like
+        _GeneratorContextManager to support use as
+        decorators via implicit recreation.
+        
+        Note: this is a private interface just for _GCM in 3.2 but will be
+        renamed and documented for third party use in 3.3
+        """
+        return self
+
     def __call__(self, func):
         @wraps(func)
         def inner(*args, **kwds):
-            with self:
+            with self._recreate_cm():
                 return func(*args, **kwds)
         return inner
 
@@ -20,8 +33,15 @@ class ContextDecorator(object):
 class _GeneratorContextManager(ContextDecorator):
     """Helper for @contextmanager decorator."""
 
-    def __init__(self, gen):
-        self.gen = gen
+    def __init__(self, func, *args, **kwds):
+        self.gen = func(*args, **kwds)
+        self.func, self.args, self.kwds = func, args, kwds
+
+    def _recreate_cm(self):
+        # _GCM instances are one-shot context managers, so the
+        # CM must be recreated each time a decorated function is
+        # called
+        return self.__class__(self.func, *self.args, **self.kwds)
 
     def __enter__(self):
         try:
@@ -92,7 +112,7 @@ def contextmanager(func):
     """
     @wraps(func)
     def helper(*args, **kwds):
-        return _GeneratorContextManager(func(*args, **kwds))
+        return _GeneratorContextManager(func, *args, **kwds)
     return helper
 
 
diff --git a/Lib/test/test_contextlib.py b/Lib/test/test_contextlib.py
index d6bb5b818e..6e38305123 100644
--- a/Lib/test/test_contextlib.py
+++ b/Lib/test/test_contextlib.py
@@ -350,13 +350,13 @@ class TestContextDecorator(unittest.TestCase):
 
 
     def test_contextmanager_as_decorator(self):
-        state = []
         @contextmanager
         def woohoo(y):
             state.append(y)
             yield
             state.append(999)
 
+        state = []
         @woohoo(1)
         def test(x):
             self.assertEqual(state, [1])
@@ -364,6 +364,11 @@ class TestContextDecorator(unittest.TestCase):
         test('something')
         self.assertEqual(state, [1, 'something', 999])
 
+        # Issue #11647: Ensure the decorated function is 'reusable'
+        state = []
+        test('something else')
+        self.assertEqual(state, [1, 'something else', 999])
+
 
 # This is needed to make the test actually run under regrtest.py!
 def test_main():
diff --git a/Lib/test/test_with.py b/Lib/test/test_with.py
index a9d374b324..e8cc8c056e 100644
--- a/Lib/test/test_with.py
+++ b/Lib/test/test_with.py
@@ -14,8 +14,8 @@ from test.support import run_unittest
 
 
 class MockContextManager(_GeneratorContextManager):
-    def __init__(self, gen):
-        _GeneratorContextManager.__init__(self, gen)
+    def __init__(self, func, *args, **kwds):
+        super().__init__(func, *args, **kwds)
         self.enter_called = False
         self.exit_called = False
         self.exit_args = None
@@ -33,7 +33,7 @@ class MockContextManager(_GeneratorContextManager):
 
 def mock_contextmanager(func):
     def helper(*args, **kwds):
-        return MockContextManager(func(*args, **kwds))
+        return MockContextManager(func, *args, **kwds)
     return helper
 
 
diff --git a/Misc/ACKS b/Misc/ACKS
index c3f4e94ba9..0443e9374c 100644
--- a/Misc/ACKS
+++ b/Misc/ACKS
@@ -704,6 +704,7 @@ Burton Radons
 Brodie Rao
 Antti Rasinen
 Sridhar Ratnakumar
+Ysj Ray
 Eric Raymond
 Edward K. Ream
 Chris Rebert
diff --git a/Misc/NEWS b/Misc/NEWS
index ca22acf2ab..5fb153d5ec 100644
--- a/Misc/NEWS
+++ b/Misc/NEWS
@@ -83,6 +83,10 @@ Core and Builtins
 Library
 -------
 
+- Issue #11647: objects created using contextlib.contextmanager now support
+  more than one call to the function when used as a decorator. Initial patch
+  by Ysj Ray.
+
 - logging: don't define QueueListener if Python has no thread support.
 
 - functools.cmp_to_key() now works with collections.Hashable().