1. 程式人生 > >實現一個重試裝飾器

實現一個重試裝飾器

class retry(object):
"""A retry decorator."""

def __init__(self,
exception=Exception,
timeout=None,
retries=None,
interval=0.001,
logfun=lambda s: print(s, file=sys.stderr),
):
if timeout and retries:
raise ValueError("timeout and retries args are mutually exclusive")
self.exception = exception
self.timeout = timeout
self.retries = retries
self.interval = interval
self.logfun = logfun

def __iter__(self):
if self.timeout:
stop_at = time.time() + self.timeout
while time.time() < stop_at:
yield
elif self.retries:
for _ in range(self.retries):
yield
else:
while True:
yield

def sleep(self):
if self.interval is not None:
time.sleep(self.interval)

def __call__(self, fun):
@functools.wraps(fun)
def wrapper(*args, **kwargs):
exc = None
for _ in self:
try:
return fun(*args, **kwargs)
except self.exception as _:
exc = _
if self.logfun is not None:
self.logfun(exc)
self.sleep()
continue
if PY3:
raise exc
else:
raise

# This way the user of the decorated function can change config
# parameters.
wrapper.decorator = self
return wrapper

def retry_before_failing(retries=NO_RETRIES):
"""Decorator which runs a test function and retries N times before
actually failing.
"""
return retry(exception=AssertionError, timeout=None, retries=retries)


@retry_before_failing()
def test():
  pass