Merge 0b35562 into f310835

deal/_decorators/base.py

@@ -1,7 +1,7 @@
 # built-in
 from asyncio import iscoroutinefunction
 from functools import update_wrapper
-from inspect import getcallargs
+from inspect import getcallargs, isgeneratorfunction
 from typing import Callable, Type
 
 # app
@@ -94,6 +94,14 @@ async def async_wrapped(*args, **kwargs):
             else:
                 return await function(*args, **kwargs)
 
+        def wrapped_generator(*args, **kwargs):
+            if self.enabled:
+                yield from self.patched_generator(*args, **kwargs)
+            else:
+                yield from function(*args, **kwargs)
+
         if iscoroutinefunction(function):
             return update_wrapper(async_wrapped, function)
+        if isgeneratorfunction(function):
+            return update_wrapper(wrapped_generator, function)
         return update_wrapper(wrapped, function)

deal/_decorators/ensure.py

@@ -26,3 +26,11 @@ async def async_patched_function(self, *args, **kwargs):
         result = await self.function(*args, **kwargs)
         self.validate(*args, result=result, **kwargs)
         return result
+
+    def patched_generator(self, *args, **kwargs):
+        """
+        Step 3. Wrapped function calling.
+        """
+        for result in self.function(*args, **kwargs):
+            self.validate(*args, result=result, **kwargs)
+            yield result

deal/_decorators/offline.py

@@ -25,23 +25,43 @@ def patched_function(self, *args, **kwargs):
         """
         Step 3. Wrapped function calling.
         """
-        true_socket = socket.socket
-        socket.socket = self.fake_socket
+        self.patch()
         try:
             return self.function(*args, **kwargs)
         finally:
-            socket.socket = true_socket
+            self.unpatch()
 
     async def async_patched_function(self, *args, **kwargs):
         """
         Step 3. Wrapped function calling.
         """
-        true_socket = socket.socket
-        socket.socket = self.fake_socket
+        self.patch()
         try:
             return await self.function(*args, **kwargs)
         finally:
-            socket.socket = true_socket
+            self.unpatch()
+
+    def patched_generator(self, *args, **kwargs):
+        """
+        Step 3. Wrapped function calling.
+        """
+        generator = self.function(*args, **kwargs)
+        while True:
+            self.patch()
+            try:
+                result = next(generator)
+            except StopIteration:
+                return
+            finally:
+                self.unpatch()
+            yield result
+
+    def patch(self):
+        self.true_socket = socket.socket
+        socket.socket = self.fake_socket
+
+    def unpatch(self):
+        socket.socket = self.true_socket
 
     def fake_socket(self, *args, **kwargs):
         raise self.exception

deal/_decorators/post.py

@@ -26,3 +26,11 @@ async def async_patched_function(self, *args, **kwargs):
         result = await self.function(*args, **kwargs)
         self.validate(result)
         return result
+
+    def patched_generator(self, *args, **kwargs):
+        """
+        Step 3. Wrapped function calling.
+        """
+        for result in self.function(*args, **kwargs):
+            self.validate(result)
+            yield result

deal/_decorators/pre.py

@@ -24,3 +24,10 @@ async def async_patched_function(self, *args, **kwargs):
         """
         self.validate(*args, **kwargs)
         return await self.function(*args, **kwargs)
+
+    def patched_generator(self, *args, **kwargs):
+        """
+        Step 3. Wrapped function calling.
+        """
+        self.validate(*args, **kwargs)
+        yield from self.function(*args, **kwargs)

deal/_decorators/raises.py

@@ -47,3 +47,16 @@ async def async_patched_function(self, *args, **kwargs):
             if not isinstance(exc, self.exceptions):
                 raise self.exception from exc
             raise
+
+    def patched_generator(self, *args, **kwargs):
+        """
+        Step 3. Wrapped function calling.
+        """
+        try:
+            yield from self.function(*args, **kwargs)
+        except ContractError:
+            raise
+        except Exception as exc:
+            if not isinstance(exc, self.exceptions):
+                raise self.exception from exc
+            raise

deal/_decorators/reason.py

@@ -48,3 +48,16 @@ async def async_patched_function(self, *args, **kwargs):
             except self.exception:
                 raise self.exception from origin
             raise
+
+    def patched_generator(self, *args, **kwargs):
+        """
+        Step 3. Wrapped function calling.
+        """
+        try:
+            yield from self.function(*args, **kwargs)
+        except self.event as origin:
+            try:
+                self.validate(*args, **kwargs)
+            except self.exception:
+                raise self.exception from origin
+            raise

deal/_decorators/silent.py

@@ -5,7 +5,7 @@
 # app
 from .._exceptions import SilentContractError
 from .._types import ExceptionType
-from .base import Base
+from .offline import Offline
 
 
 class PatchedStringIO(StringIO):
@@ -16,44 +16,15 @@ def write(self, *args, **kwargs):
         raise self.exception
 
 
-class Silent(Base):
+class Silent(Offline):
     exception: ExceptionType = SilentContractError
 
-    def __init__(self, *, message: str = None, exception: ExceptionType = None, debug: bool = False):
-        """
-        Step 1. Init params.
-        """
-        super().__init__(
-            validator=None,
-            message=message,
-            exception=exception,
-            debug=debug,
-        )
-
-    def patched_function(self, *args, **kwargs):
-        """
-        Step 3. Wrapped function calling.
-        """
-        true_stdout = sys.stdout
-        true_stderr = sys.stderr
+    def patch(self):
+        self.true_stdout = sys.stdout
+        self.true_stderr = sys.stderr
         sys.stdout = PatchedStringIO(exception=self.exception)
         sys.stderr = PatchedStringIO(exception=self.exception)
-        try:
-            return self.function(*args, **kwargs)
-        finally:
-            sys.stdout = true_stdout
-            sys.stderr = true_stderr
-
-    async def async_patched_function(self, *args, **kwargs):
-        """
-        Step 3. Wrapped function calling.
-        """
-        true_stdout = sys.stdout
-        true_stderr = sys.stderr
-        sys.stdout = PatchedStringIO(exception=self.exception)
-        sys.stderr = PatchedStringIO(exception=self.exception)
-        try:
-            return await self.function(*args, **kwargs)
-        finally:
-            sys.stdout = true_stdout
-            sys.stderr = true_stderr
+
+    def unpatch(self):
+        sys.stdout = self.true_stdout
+        sys.stderr = self.true_stderr

tests/test_decorators/test_ensure.py

@@ -32,3 +32,15 @@ async def func(a, b):
     assert run_sync(func(1, 2)) == 'different numbers'
     with pytest.raises(deal.PostContractError):
         run_sync(func(0, 1))
+
+
+def test_decorating_generator():
+    @deal.ensure(lambda x, y, result: result > y ** 2)
+    def multiply(x, y):
+        yield x * y
+        yield x * y * 2
+        yield x * y * 4
+
+    assert list(multiply(2, 1)) == [2, 4, 8]
+    with pytest.raises(deal.PostContractError):
+        list(multiply(2, 2))

tests/test_decorators/test_offline.py

@@ -42,3 +42,17 @@ async def func(do):
     assert run_sync(func(False)) == 1
     with pytest.raises(deal.OfflineContractError):
         run_sync(func(True))
+
+
+def test_decorating_generator():
+    @deal.offline
+    def func(do):
+        if not do:
+            yield 1
+            return
+        http = urllib3.PoolManager()
+        http.request('GET', 'http://httpbin.org/robots.txt')
+
+    assert list(func(False)) == [1]
+    with pytest.raises(deal.OfflineContractError):
+        list(func(True))

tests/test_decorators/test_post.py

@@ -20,3 +20,15 @@ async def func(x):
     assert run_sync(func(-2)) == 2
     with pytest.raises(deal.PostContractError):
         run_sync(func(2))
+
+
+def test_decorating_generator():
+    @deal.post(lambda x: x <= 8)
+    def double(x):
+        yield x
+        yield x * 2
+        yield x * 4
+
+    assert list(double(2)) == [2, 4, 8]
+    with pytest.raises(deal.PostContractError):
+        list(double(4))

tests/test_decorators/test_pre.py

@@ -129,3 +129,15 @@ async def double(x):
     assert run_sync(double(2)) == 4
     with pytest.raises(deal.PreContractError):
         run_sync(double(-2))
+
+
+def test_decorating_generator():
+    @deal.pre(lambda x: x > 0)
+    def double(x):
+        yield x
+        yield x * 2
+        yield x * 4
+
+    assert list(double(2)) == [2, 4, 8]
+    with pytest.raises(deal.PreContractError):
+        list(double(-2))

tests/test_decorators/test_raises.py

@@ -60,3 +60,33 @@ async def func(do, number):
         run_sync(func(True, 1))
     with pytest.raises(ZeroDivisionError):
         run_sync(func(False, 0))
+
+
+def test_decorating_generator():
+    @deal.raises(ZeroDivisionError)
+    def func(x):
+        if x == -1:
+            raise KeyError
+        yield 1 / x
+
+    assert list(func(1)) == [1]
+    with pytest.raises(ZeroDivisionError):
+        list(func(0))
+    with pytest.raises(deal.RaisesContractError):
+        list(func(-1))
+
+
+def test_raises_generator():
+    @deal.raises(ZeroDivisionError)
+    @deal.offline
+    def func(do, number):
+        if do:
+            http = urllib3.PoolManager()
+            http.request('GET', 'http://httpbin.org/robots.txt')
+        yield 1 / number
+
+    assert list(func(False, 1)) == [1]
+    with pytest.raises(deal.OfflineContractError):
+        list(func(True, 1))
+    with pytest.raises(ZeroDivisionError):
+        list(func(False, 0))

tests/test_decorators/test_reason.py

@@ -48,3 +48,24 @@ async def func(x):
     assert run_sync(func(3)) == 3
     with pytest.raises(exc):
         run_sync(func(value))
+
+
+@pytest.mark.parametrize('value, exc', [
+    (0, ZeroDivisionError),
+    (1, ValueError),
+    (2, deal.ReasonContractError),
+])
+def test_decorating_generator(value, exc):
+    @deal.reason(ValueError, lambda x: x == 1)
+    def func(x):
+        if x == 0:
+            raise ZeroDivisionError
+        if x == 1:
+            raise ValueError
+        if x == 2:
+            raise ValueError
+        yield x
+
+    assert list(func(3)) == [3]
+    with pytest.raises(exc):
+        list(func(value))

tests/test_decorators/test_silent.py

@@ -25,3 +25,15 @@ async def func(msg):
     assert run_sync(func('')) == ''
     with pytest.raises(deal.SilentContractError):
         run_sync(func('a'))
+
+
+def test_decorating_generator():
+    @deal.silent
+    def func(msg):
+        if msg:
+            print(msg)
+        yield msg
+
+    assert list(func('')) == ['']
+    with pytest.raises(deal.SilentContractError):
+        list(func('a'))

tests/test_state.py

@@ -32,3 +32,15 @@ async def func(x):
     deal.switch(main=True)
     with pytest.raises(deal.PreContractError):
         run_sync(func(-2))
+
+
+def test_contract_state_switch_default_param_generator():
+    @deal.pre(lambda x: x > 0)
+    def func(x):
+        yield x * 2
+
+    deal.switch(main=False)
+    assert list(func(-2)) == [-4]
+    deal.switch(main=True)
+    with pytest.raises(deal.PreContractError):
+        list(func(-2))