Baaa

cool-RR · cool-RR · commit 79456e58f85a · 2015-12-25T16:14:19.000+02:00
diff --git a/source_py3/python_toolbox/context_management/functions.py b/source_py3/python_toolbox/context_management/functions.py
@@ -8,6 +8,7 @@
 '''
 
 import sys
+import types
 
 from python_toolbox import caching
 
@@ -45,6 +46,15 @@ def nested(*managers):
             raise exc[1].with_traceback(exc[2])
     
 
+def _wrap_context_manager_or_context_generator(thing, wrapper_factory):
+    if hasattr(type(thing), '__enter__'):
+        # It's a context manager.
+        return wrapper_factory(thing)
+    else:
+        assert isinstance(thing, types.FunctionType)
+        # It's a context generator.
+        
+
 
 def as_idempotent(context_manager):
     '''
diff --git a/source_py3/test_python_toolbox/test_context_management/test_as_reentrant.py b/source_py3/test_python_toolbox/test_context_management/test_as_reentrant.py
@@ -3,7 +3,8 @@
 
 import queue as queue_module
 
-from python_toolbox.context_management import as_reentrant, ContextManager
+from python_toolbox.context_management import (as_reentrant, ContextManager,
+                                               ContextManagerType)
 from python_toolbox import cute_testing
 
 
@@ -149,7 +150,7 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
     # happens *after* `.__wrapped__.__enter__` is called:
     assert depth_log.get(block=False) == 1
     
-def test_decorator():
+def test_decorator_manage_context():
     class Meow(ContextManager):
         n = 0
             
@@ -175,4 +176,81 @@ def manage_context(self):
     assert meow.n == 0
         
         
+def test_decorator_class():
+    
+    @as_reentrant
+    class Meow(ContextManager):
+        n = 0
+            
+        def manage_context(self):
+            self.n += 1
+            try:
+                yield
+            finally:
+                self.n -= 1
+                
+            
+    meow = Meow()
+    assert meow.n == 0
+    with meow:
+        assert meow.n == 1
+        with meow:
+            assert meow.n == 1
+            with meow:
+                assert meow.n == 1
+            assert meow.n == 1
+        assert meow.n == 1
+    assert meow.n == 0
+        
+def test_decorator_class_enter_exit():
+    
+    @as_reentrant
+    class Meow(ContextManager):
+        n = 0
+            
+        def __enter__(self):
+            self.n += 1
+            return self
+        
+        def __exit__(self, exc_type, exc_value, exc_traceback):
+            self.n -= 1
+                
+            
+    meow = Meow()
+    assert meow.n == 0
+    with meow:
+        assert meow.n == 1
+        with meow:
+            assert meow.n == 1
+            with meow:
+                assert meow.n == 1
+            assert meow.n == 1
+        assert meow.n == 1
+    assert meow.n == 0
+        
+        
+def test_decorator_decorator():
+    @as_reentrant
+    @ContextManagerType
+    def Meow():
+        Meow.n += 1
+        try:
+            yield
+        finally:
+            Meow.n -= 1
+                
+            
+    meow = Meow()
+    assert Meow.n == 0
+    with meow:
+        assert Meow.n == 1
+        with meow:
+            assert Meow.n == 1
+            with meow:
+                assert Meow.n == 1
+            assert Meow.n == 1
+        assert Meow.n == 1
+    assert Meow.n == 0
+        
+        
         
