Merge pull request #124 from immerrr/make-before-sleep-accept-call-state

Make before_sleep accept call_state

Author: mergify[bot]

.gitignore

@@ -8,3 +8,5 @@ AUTHORS
 ChangeLog
 .eggs/
 doc/_build
+
+/.pytest_cache

doc/source/conf.py

@@ -22,7 +22,8 @@
 project = "Tenacity"
 
 extensions = [
-    'sphinx.ext.doctest'
+    'sphinx.ext.doctest',
+    'sphinx.ext.autodoc',
 ]
 
 # -- Options for sphinx.ext.doctest  -----------------------------------------

doc/source/index.rst

@@ -272,6 +272,109 @@ In the same spirit, It's possible to execute after a call that failed:
     def raise_my_exception():
         raise MyException("Fail")
 
+It's also possible to only log failures that are going to be retried. Normally
+retries happen after a wait interval, so the keyword argument is called
+``before_sleep``:
+
+.. testcode::
+
+    logger = logging.getLogger(__name__)
+    @retry(stop=stop_after_attempt(3),
+           before_sleep=before_sleep_log(logger, logging.DEBUG))
+    def raise_my_exception():
+        raise MyException("Fail")
+
+You can also define a custom ``before_sleep`` function. It should have one
+parameter called ``call_state`` that contains all information about current
+retry invocation.
+
+.. testcode::
+
+    logger = logging.getLogger(__name__)
+
+    def my_before_sleep(call_state):
+        if call_state.attempt_number < 1:
+            loglevel = logging.INFO
+        else:
+            loglevel = logging.WARNING
+        logger.log(
+            loglevel, 'Retrying %s: attempt %s ended with: %s',
+            call_state.fn, call_state.attempt_number, call_state.outcome)
+
+    @retry(stop=stop_after_attempt(3), before_sleep=my_before_sleep)
+    def raise_my_exception():
+        raise MyException("Fail")
+
+    try:
+        raise_my_exception()
+    except RetryError:
+        pass
+
+``call_state`` argument is an object of `RetryCallState` class:
+
+.. autoclass:: tenacity.RetryCallState
+
+   Constant attributes:
+
+   .. autoinstanceattribute:: start_time(float)
+      :annotation:
+
+   .. autoinstanceattribute:: retry_object(BaseRetrying)
+      :annotation:
+
+   .. autoinstanceattribute:: fn(callable)
+      :annotation:
+
+   .. autoinstanceattribute:: args(tuple)
+      :annotation:
+
+   .. autoinstanceattribute:: kwargs(dict)
+      :annotation:
+
+   Variable attributes:
+
+   .. autoinstanceattribute:: attempt_number(int)
+      :annotation:
+
+   .. autoinstanceattribute:: outcome(tenacity.Future or None)
+      :annotation:
+
+   .. autoinstanceattribute:: outcome_timestamp(float or None)
+      :annotation:
+
+   .. autoinstanceattribute:: idle_for(float)
+      :annotation:
+
+   .. autoinstanceattribute:: next_action(tenacity.RetryAction or None)
+      :annotation:
+
+
+
+Previously custom ``before_sleep`` functions could also accept three
+parameters. This approach is deprecated but kept for backward compatibility:
+
+- ``retry_object``: the `Retrying` object that runs the retries,
+- ``sleep``: wait interval in seconds before the next attempt,
+- ``last_result``: the outcome of the current attempt.
+
+.. testcode::
+
+    logger = logging.getLogger(__name__)
+
+    def my_before_sleep(retry_object, sleep, last_result):
+        logger.warning(
+            'Retrying %s: last_result=%s, retrying in %s seconds...',
+            retry_object.fn, last_result, sleep)
+
+    @retry(stop=stop_after_attempt(3), before_sleep=my_before_sleep)
+    def raise_my_exception():
+        raise MyException("Fail")
+
+    try:
+        raise_my_exception()
+    except RetryError:
+        pass
+
 Similarly, you can call a custom callback function after all retries failed, without raising an exception (or you can re-raise or do anything really)
 
 .. testcode::

tenacity/__init__.py

@@ -34,6 +34,7 @@
 import six
 
 from tenacity import _utils
+from tenacity import before_sleep as _before_sleep
 from tenacity import wait as _wait
 
 # Import all built-in retry strategies for easier usage.
@@ -122,6 +123,35 @@ class DoSleep(float):
     pass
 
 
+class BaseAction(object):
+    """Base class for representing actions to take by retry object.
+
+    Concrete implementations must define:
+    - __init__: to initialize all necessary fields
+    - REPR_ATTRS: class variable specifying attributes to include in repr(self)
+    - NAME: for identification in retry object methods and callbacks
+    """
+
+    REPR_FIELDS = ()
+    NAME = None
+
+    def __repr__(self):
+        state_str = ', '.join('%s=%r' % (field, getattr(self, field))
+                              for field in self.REPR_FIELDS)
+        return '%s(%s)' % (type(self).__name__, state_str)
+
+    def __str__(self):
+        return repr(self)
+
+
+class RetryAction(BaseAction):
+    REPR_FIELDS = ('sleep',)
+    NAME = 'retry'
+
+    def __init__(self, sleep):
+        self.sleep = float(sleep)
+
+
 _unset = object()
 
 
@@ -148,7 +178,7 @@ def __init__(self,
                  retry=retry_if_exception_type(),
                  before=before_nothing,
                  after=after_nothing,
-                 before_sleep=before_sleep_nothing,
+                 before_sleep=None,
                  reraise=False,
                  retry_error_cls=RetryError,
                  retry_error_callback=None):
@@ -158,16 +188,25 @@ def __init__(self,
         self.retry = retry
         self.before = before
         self.after = after
-        self.before_sleep = before_sleep
+        self._before_sleep = before_sleep
         self.reraise = reraise
         self._local = threading.local()
         self.retry_error_cls = retry_error_cls
         self.retry_error_callback = retry_error_callback
 
+        # This attribute was moved to RetryCallState and is deprecated on
+        # Retrying objects but kept for backward compatibility.
+        self.fn = None
+
     @_utils.cached_property
     def wait(self):
         return _wait._wait_func_accept_call_state(self._wait)
 
+    @_utils.cached_property
+    def before_sleep(self):
+        return _before_sleep._before_sleep_func_accept_call_state(
+            self._before_sleep)
+
     def copy(self, sleep=_unset, stop=_unset, wait=_unset,
              retry=_unset, before=_unset, after=_unset, before_sleep=_unset,
              reraise=_unset):
@@ -244,6 +283,7 @@ def begin(self, fn):
         self.statistics['start_time'] = _utils.now()
         self.statistics['attempt_number'] = 1
         self.statistics['idle_for'] = 0
+        self.fn = fn
 
     def iter(self, call_state):  # noqa
         fut = call_state.outcome
@@ -275,13 +315,14 @@ def iter(self, call_state):  # noqa
         if self.wait:
             sleep = self.wait(call_state=call_state)
         else:
-            sleep = 0
+            sleep = 0.0
+        call_state.next_action = RetryAction(sleep)
         call_state.idle_for += sleep
         self.statistics['idle_for'] += sleep
         self.statistics['attempt_number'] += 1
 
         if self.before_sleep is not None:
-            self.before_sleep(self, sleep=sleep, last_result=fut)
+            self.before_sleep(call_state=call_state)
 
         return DoSleep(sleep)
 
@@ -292,7 +333,8 @@ class Retrying(BaseRetrying):
     def call(self, fn, *args, **kwargs):
         self.begin(fn)
 
-        call_state = RetryCallState(fn=fn, args=args, kwargs=kwargs)
+        call_state = RetryCallState(
+            retry_object=self, fn=fn, args=args, kwargs=kwargs)
         while True:
             do = self.iter(call_state=call_state)
             if isinstance(do, DoAttempt):
@@ -337,16 +379,28 @@ def construct(cls, attempt_number, value, has_exception):
 class RetryCallState(object):
     """State related to a single call wrapped with Retrying."""
 
-    def __init__(self, fn, args, kwargs):
+    def __init__(self, retry_object, fn, args, kwargs):
+        #: Retry call start timestamp
         self.start_time = _utils.now()
+        #: Retry manager object
+        self.retry_object = retry_object
+        #: Function wrapped by this retry call
         self.fn = fn
+        #: Arguments of the function wrapped by this retry call
         self.args = args
+        #: Keyword arguments of the function wrapped by this retry call
         self.kwargs = kwargs
 
-        self.idle_for = 0
+        #: The number of the current attempt
+        self.attempt_number = 1
+        #: Last outcome (result or exception) produced by the function
         self.outcome = None
+        #: Timestamp of the last outcome
         self.outcome_timestamp = None
-        self.attempt_number = 1
+        #: Time spent sleeping in retries
+        self.idle_for = 0
+        #: Next action as decided by the retry manager
+        self.next_action = None
 
     @property
     def seconds_since_start(self):
@@ -358,6 +412,7 @@ def prepare_for_next_attempt(self):
         self.outcome = None
         self.outcome_timestamp = None
         self.attempt_number += 1
+        self.next_action = None
 
     def set_result(self, val):
         ts = _utils.now()

tenacity/_asyncio.py

@@ -37,7 +37,8 @@ def __init__(self,
     def call(self, fn, *args, **kwargs):
         self.begin(fn)
 
-        call_state = RetryCallState(fn=fn, args=args, kwargs=kwargs)
+        call_state = RetryCallState(
+            retry_object=self, fn=fn, args=args, kwargs=kwargs)
         while True:
             do = self.iter(call_state=call_state)
             if isinstance(do, DoAttempt):

tenacity/_utils.py

@@ -132,3 +132,13 @@ def __get__(self, obj, cls):
             return self
         value = obj.__dict__[self.func.__name__] = self.func(obj)
         return value
+
+
+def _func_takes_call_state(func):
+    if not six.callable(func):
+        return False
+    if not inspect.isfunction(func):
+        # func is a callable object rather than a function
+        func = func.__call__
+    func_spec = getargspec(func)
+    return 'call_state' in func_spec.args

tenacity/before_sleep.py

@@ -14,20 +14,43 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import six
+
 from tenacity import _utils
 
 
-def before_sleep_nothing(retry_obj, sleep, last_result):
+def before_sleep_nothing(call_state, sleep, last_result):
     """Before call strategy that does nothing."""
 
 
 def before_sleep_log(logger, log_level):
     """Before call strategy that logs to some logger the attempt."""
-    def log_it(retry_obj, sleep, last_result):
-        logger.log(log_level,
-                   "Retrying %s in %d seconds as it raised %s.",
-                   _utils.get_callback_name(retry_obj.fn),
-                   sleep,
-                   last_result.exception())
+    def log_it(call_state):
+        if call_state.outcome.failed:
+            verb, value = 'raised', call_state.outcome.exception()
+        else:
+            verb, value = 'returned', call_state.outcome.result()
 
+        logger.log(log_level,
+                   "Retrying %s in %s seconds as it %s %s.",
+                   _utils.get_callback_name(call_state.fn),
+                   getattr(call_state.next_action, 'sleep'),
+                   verb, value)
     return log_it
+
+
+def _before_sleep_func_accept_call_state(fn):
+    if not six.callable(fn):
+        return fn
+
+    takes_call_state = _utils._func_takes_call_state(fn)
+    if takes_call_state:
+        return fn
+
+    @six.wraps(fn)
+    def wrapped_before_sleep_func(call_state):
+        return fn(
+            call_state.retry_object,
+            sleep=getattr(call_state.next_action, 'sleep'),
+            last_result=call_state.outcome)
+    return wrapped_before_sleep_func