01 марта 2017 г. в 23:44
Python: декоратор с опциональной передачей аргументов

В процессе разработки нашего децентрализованного мультипротокольного мессенджера возникла задача создать декоратор для функции, способный вызываться как в виде @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)