В процессе разработки нашего децентрализованного мультипротокольного мессенджера возникла задача создать декоратор для функции, способный вызываться как в виде @decorator так и @decorator(arg1, arg2). Мой товарищ MrBoriska создал сниппет, решающий эту проблему:
from functools import wraps
import inspect
def decorator_gen(*args, **kwargs):
# Выясняем в каком режиме работаем
as_decorator = False
# Если передан один позиционный аргумент и он является callable обьектом
if len(args) == 1 and len(kwargs) == 0 and callable(args[0]):
# Необходимо дополнительно проверить, декорирован ли он нашим декоратором
for line in inspect.getsourcelines(args[0])[0]:
line = line.strip()
if line.startswith("@decorator_gen"): # для параноиков можно еще проверять на отсутствие скобок
as_decorator = True
break
if line.startswith("def") or line.startswith("async def"):
break
# >> причесываем args и kwargs тут <<
def decorator(f):
print("i am decorator")
@wraps(f)
def wrapper(*af, **kwf):
# af и kwf - аргументы декорируемой функции f
# args и kwargs - дополнительные параметры декоратора
print("i am wrapper", (af,kwf))
if as_decorator:
print("no \"()\", no special args")
f(*af, **kwf)
return wrapper
if as_decorator:
# Если первый позиционный аргумент - функция, то работаем как декоратор
return decorator(args[0])
else:
# Иначе, конфигурируем декоратор, и возвращаем его обьект-функцию
print("i am generate decorator with arguments:", (args, kwargs))
return decorator
@decorator_gen() # equal to @decorator_gen
def func(*args,**kwargs):
print("i am func",(args,kwargs))
func(1,2)Такой дескриптор будет железобетонно определять, какие параметры переданы и соответственно отрабатывать, однако если не предполагается, что первым (и единственным) аргументом в дескриптор может может передана callable функция, то можно упростить проверки и не подключать библиотеку inspect:
from functools import wraps
def decorator_gen(*args, **kwargs):
# Выясняем в каком режиме работаем
as_decorator = len(args) == 1 and len(kwargs) == 0 and callable(args[0])
# >> причесываем args и kwargs тут <<
def decorator(f):
print("i am decorator")
@wraps(f)
def wrapper(*af, **kwf):
# af и kwf - аргументы декорируемой функции f
# args и kwargs - дополнительные параметры декоратора
print("i am wrapper", (af,kwf))
if as_decorator:
print("no \"()\", no special args")
f(*af, **kwf)
return wrapper
if as_decorator:
# Если первый позиционный аргумент - функция, то работаем как декоратор
return decorator(args[0])
else:
# Иначе, конфигурируем декоратор, и возвращаем его обьект-функцию
print("i am generate decorator with arguments:", (args, kwargs))
return decorator
@decorator_gen() # equal to @decorator_gen
def func(*args,**kwargs):
print("i am func",(args,kwargs))
func(1,2)А для того, чтобы подробнее узнать о работе с декораторами в Python, рекомендую почитать по этим ссылкам: devacademy (ru) и Python 3 Patterns, Recipes and Idioms (en)