-

cool-RR · cool-RR · commit f24be3521d64 · 2015-12-12T00:42:40.000+02:00
diff --git a/source_py3/python_toolbox/context_management/__init__.py b/source_py3/python_toolbox/context_management/__init__.py
@@ -133,4 +133,4 @@ def do_stuff():
 from .blank_context_manager import BlankContextManager
 from .reentrant_context_manager import ReentrantContextManager
 from .delegating_context_manager import DelegatingContextManager
-from .functions import nested
+from .functions import nested, idempotentify
diff --git a/source_py3/python_toolbox/context_management/functions.py b/source_py3/python_toolbox/context_management/functions.py
@@ -10,6 +10,7 @@
 import sys
 
 from .context_manager_type import ContextManagerType
+from .context_manager import ContextManager
 
 
 @ContextManagerType
@@ -40,4 +41,50 @@ def nested(*managers):
             # the right information. Another exception may
             # have been raised and caught by an exit method
             raise exc[1].with_traceback(exc[2])
-    
+    
+    
+def idempotentify(context_manager):
+    '''
+    Wrap a context manager so repeated calls to enter and exit will be ignored.
+    
+    This means that if you call `__enter__` a second time on the context
+    manager, nothing will happen. The `__enter__` method won't be called and an
+    exception would not be raised. Same goes for the `__exit__` method, after
+    calling it once, if you try to call it again it will be a no-op. But now
+    that you've called `__exit__` you can call `__enter__` and it will really
+    do the enter action again, and then `__exit__` will be available again,
+    etc.
+    
+    This is useful when you have a context manager that you want to put in an
+    `ExitStack`, but you also possibly want to exit it manually before the
+    `ExitStack` closes. This way you don't risk an exception by having the
+    context manager exit twice.
+    '''
+    class _IdempotentContextManager:
+        # Not inheriting from `ContextManager` because this class is internal
+        # and not used beyond `idempotentify`, so no need for the extra
+        # baggage.
+        _entered = False
+        _enter_value = None
+        
+        def __init__(self, wrapped_context_manager):
+            self.__wrapped__ = wrapped_context_manager
+            
+        
+        def __enter__(self):
+            if not self._entered:
+                self._enter_value = self.__wrapped__.__enter__()
+                self._entered = True
+            return self._enter_value
+                
+            
+        def __exit__(self, exc_type, exc_value, exc_traceback):
+            if self._entered:
+                exit_value = self.__wrapped__.__exit__(exc_type, exc_value,
+                                                       exc_traceback)
+                self._entered = False
+                return exit_value
+                
+    return _IdempotentContextManager(context_manager)
+    
+            
diff --git a/source_py3/test_python_toolbox/test_context_management/test_idempotentify.py b/source_py3/test_python_toolbox/test_context_management/test_idempotentify.py
@@ -0,0 +1,62 @@
+# Copyright 2009-2015 Ram Rachum.
+# This program is distributed under the MIT license.
+
+'''Testing module for `context_management.ReentrantContextManager`.''' 
+
+import queue as queue_module
+
+from python_toolbox.context_management import idempotentify, ContextManager
+from python_toolbox import cute_testing
+
+
+class SomeContextManager(ContextManager):
+    x = 0
+    def manage_context(self):
+        self.x += 1
+        try:
+            yield self
+        finally:
+            self.x -= 1
+            
+            
+
+def test_idempotentify():
+    some_context_manager = SomeContextManager()
+    assert some_context_manager.x == 0
+    with some_context_manager:
+        assert some_context_manager.x == 1
+    assert some_context_manager.x == 0
+    
+    some_context_manager.__enter__()
+    assert some_context_manager.x == 1
+    some_context_manager.__enter__()
+    assert some_context_manager.x == 2
+    some_context_manager.__enter__()
+    assert some_context_manager.x == 3
+    some_context_manager.__exit__(None, None, None)
+    assert some_context_manager.x == 2
+    some_context_manager.__exit__(None, None, None)
+    assert some_context_manager.x == 1
+    some_context_manager.__exit__(None, None, None)
+    assert some_context_manager.x == 0
+    with cute_testing.RaiseAssertor():
+        some_context_manager.__exit__(None, None, None)
+    with cute_testing.RaiseAssertor():
+        some_context_manager.__exit__(None, None, None)
+        
+    idempotent_context_manager = idempotentify(SomeContextManager())
+
+    idempotent_context_manager.__enter__()
+    assert idempotent_context_manager.__wrapped__.x == 1
+    idempotent_context_manager.__enter__()
+    assert idempotent_context_manager.__wrapped__.x == 1
+    idempotent_context_manager.__enter__()
+    assert idempotent_context_manager.__wrapped__.x == 1
+    idempotent_context_manager.__exit__(None, None, None)
+    assert idempotent_context_manager.__wrapped__.x == 0
+    idempotent_context_manager.__exit__(None, None, None)
+    assert idempotent_context_manager.__wrapped__.x == 0
+    idempotent_context_manager.__exit__(None, None, None)
+    assert idempotent_context_manager.__wrapped__.x == 0
+    
+        
