之前我们介绍过了无参数和有参数装饰器，那么是否装饰器可以具备可选参数，即既可以使用@func，也可以使用@func(a=1, b=2)，答案是可以的，本篇文章我们介绍一下如何构建一个可选参数的装饰器。

我们所实现的装饰器是一个任务重启装饰器，即任务失败时，能根据指定次数和间隔时间重新运行函数。

1. 借助args

第一种方式是借助于args参数，即通过判断其长度与值的类型，在内部重新定义retry_times和sleep_time变量：

  1. import time
  2. import logging
  3. from functools import wraps
  6. def retry(*args):
  7.     def func_retry(func):
  8.         log = logging.getLogger("Retry")
  10.         @wraps(func)
  11.         def wrapper(*args, **kwargs):
  12.             cnt = 1
  14.             while cnt <= retry_times:
  15.                 try:
  16.                     res = func(*args, **kwargs)
  17.                     return res
  19.                 except Exception as e:
  20.                     log.error(f"Retry [red]{cnt}[/] times \n", extra={"markup": True})
  21.                     cnt += 1
  22.                     time.sleep(sleep_time)
  23.                     log.error("Messages: [red]{0}[/]".format(e), extra={"markup": True})
  24.             else:
  25.                 log.exception(f"Retry {retry_times} times, but all failed!")
  26.                 raise Exception(f"Retry {retry_times} times, but all failed!")
  28.         return wrapper
  30.     log = logging.getLogger("Retry")
  32.     # 对应@retry
  33.     if len(args) == 1 and callable(args[0]):
  34.         retry_times, sleep_time = 5, 60
  35.         return func_retry(args[0])
  37.     # 对应@retry(3)
  38.     elif len(args) == 1 and isinstance(args[0], int):
  39.         retry_times, sleep_time = args[0], 60
  40.         return func_retry
  42.     # 对应@retry(2, 3)
  43.     else:
  44.         retry_times, sleep_time = args[0], args[1]
  45.         return func_retry

按照上述方式所实现的装饰器，使用方法如下：

  1. @retry
  2. def retry_func():
  3.     print(1 / 0)
  5. @retry(2)
  6. def retry_func():
  7.     print(1 / 0)
  9. @retry(2, 3)
  10. def retry_func():
  11.     print(1 / 0)

可以看到，这种实现方式并不好，它并没有采用传参的方式，并且使用者对所传入值代表含义也会有疑惑。

2. 更简洁的实现方式

接下来我们介绍一种更好的实现方式，借助于functools模块中的partial函数：

  1. import time
  2. import logging
  3. from functools import wraps, partial
  5. def retry(func = None, *, retry_times: int = 5, sleep_time: int = 60):
  6.     log = logging.getLogger("Retry")
  8.     if func is None:
  9.         return partial(retry, retry_times=retry_times, sleep_time=sleep_time)
  11.     @wraps(func)
  12.     def wrapper(*args, **kwargs):
  13.         cnt = 1
  15.         while cnt <= retry_times:
  16.             try:
  17.                 res = func(*args, **kwargs)
  18.                 return res
  20.             except Exception as e:
  21.                 log.error(f"Retry [red]{cnt}[/] times \n", extra={"markup": True})
  22.                 cnt += 1
  23.                 time.sleep(sleep_time)
  24.                 log.error("Messages: [red]{0}[/]".format(e), extra={"markup": True})
  25.         else:
  26.             log.exception(f"Retry {retry_times} times, but all failed!")
  27.             raise Exception(f"Retry {retry_times} times, but all failed!")
  29.     return wrapper

显然，这种方式实现的装饰器代码简洁多了，具体使用方法如下：

  1. @retry
  2. def retry_func():
  3.     print(1 / 0)
  5. @retry(retry_times=2)
  6. def retry_func():
  7.     print(1 / 0)
  9. @retry(retry_times=2, sleep_time=3)
  10. def retry_func():
  11.     print(1 / 0)

3. 参考资料