From de3d18673097646f54b1efad5114adc7236f1060 Mon Sep 17 00:00:00 2001
From: Oleh Kuchuk <oleh.kuchuk@globallogic.com>
Date: Thu, 9 Nov 2017 18:12:55 +0200
Subject: [PATCH 1/2] Added ability to use retry as context manager

---
 async_retrying.py             | 25 ++++++++++++++++++++++++-
 tests/test_context_manager.py | 29 +++++++++++++++++++++++++++++
 2 files changed, 53 insertions(+), 1 deletion(-)
 create mode 100644 tests/test_context_manager.py

diff --git a/async_retrying.py b/async_retrying.py
index 0aa0ecc..b44b5cb 100644
--- a/async_retrying.py
+++ b/async_retrying.py
@@ -49,7 +49,7 @@ def callback(attempt, exc, args, kwargs, delay=None, *, loop):
 callback.delay = 0.5
 
 
-def retry(
+def _retry(
     fn=None,
     *,
     attempts=3,
@@ -193,3 +193,26 @@ def wrapped(*fn_args, **fn_kwargs):
         return wrapper(fn)
 
     raise NotImplementedError
+
+
+class retry(object):
+
+    def __init__(self, *args, **kwargs):
+        self._wrapper = _retry(*args, **kwargs)
+
+    def __enter__(self):
+        return self._wrapper
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    @asyncio.coroutine
+    def __aenter__(self):
+        return self._wrapper
+
+    @asyncio.coroutine
+    def __aexit__(self, exc_type, exc_val, exc_tb):
+        pass
+
+    def __call__(self, fn):
+        return self._wrapper(fn)
diff --git a/tests/test_context_manager.py b/tests/test_context_manager.py
new file mode 100644
index 0000000..5c8e6f3
--- /dev/null
+++ b/tests/test_context_manager.py
@@ -0,0 +1,29 @@
+import asyncio
+from functools import partial
+
+import pytest
+
+from async_retrying import retry
+
+@pytest.mark.run_loop
+@asyncio.coroutine
+def test_context_manager(loop):
+    counter = 0
+
+    @asyncio.coroutine
+    def fn():
+        nonlocal counter
+
+        counter += 1
+
+        if counter == 1:
+            raise RuntimeError
+
+        return counter
+
+    with retry(fn, loop=loop) as context:
+        yield from context()
+
+    ret = yield from partial(fn)()
+
+    assert ret == counter

From 3e1463c85b5251496f2d71f9502790656117fecd Mon Sep 17 00:00:00 2001
From: Oleh Kuchuk <oleh.kuchuk@globallogic.com>
Date: Fri, 10 Nov 2017 13:41:18 +0200
Subject: [PATCH 2/2] Fixed styling problems, minor usage improvements

---
 async_retrying.py | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/async_retrying.py b/async_retrying.py
index 44f63c7..910a829 100644
--- a/async_retrying.py
+++ b/async_retrying.py
@@ -49,7 +49,7 @@ def callback(attempt, exc, args, kwargs, delay=None, *, loop):
 callback.delay = 0.5
 
 
-def _retry(
+def factory(
     fn=None,
     *,
     attempts=3,
@@ -195,10 +195,10 @@ def wrapped(*fn_args, **fn_kwargs):
     raise NotImplementedError
 
 
-class retry(object):
-
-    def __init__(self, *args, **kwargs):
-        self._wrapper = _retry(*args, **kwargs)
+class retry:
+    def __init__(self, fn=None, *args, **kwargs):
+        self._fn = fn
+        self._wrapper = factory(self._fn, *args, **kwargs)
 
     def __enter__(self):
         return self._wrapper
@@ -212,7 +212,10 @@ def __aenter__(self):
 
     @asyncio.coroutine
     def __aexit__(self, exc_type, exc_val, exc_tb):
-        pass
+        return self.__exit__
+
+    def __call__(self, fn=None):
+        if not self._fn and callable(fn):
+            return self._wrapper(fn)
 
-    def __call__(self, fn):
-        return self._wrapper(fn)
+        return self._wrapper()