Python 装饰器使用的一个痛点

253 阅读2分钟

小知识,大挑战!本文正在参与“程序员必备小知识”创作活动。

本文已参与 「掘力星计划」 ,赢取创作大礼包,挑战创作激励金

如何给装饰器的参数传参,这个问题曾经困扰我好久,虽然Python版本的更新,现在这个问题终于解决了,特此记录。

疑问

首先我有一个这样的装饰器文件路径helper/log_helper.py

import traceback
from functools import wraps

from loguru import logger


def my_logger(count):
    def step1(foo):
        @wraps(foo)
        def step2(*args, **kwargs):
            try:
                result = foo(*args, **kwargs)
                logger.info(f"{result=},{count=}")
            except Exception:
                logger.exception(traceback.format_exc())

        return step2

    return step1

然后我有个文件需要引用这个装饰器demo.py

from helper.log_helper import my_logger


class Demo:
    @my_logger(count=2)
    def main(self):
        return "in main function"


if __name__ == '__main__':
    d = Demo()
    d.main()

输出结果如下

2020-10-16 11:43:12.001 | INFO     | helper.log_helper:step2:18 - result='in main function',count=2

这个装饰器的作用很简单,就是获取当前函数的返回值,和传入的count值。

好,现在问题来了?

如果给装饰器的参数传值呢,也就是说我的count=2,是通过传值的形式。你想到可能是这样

from helper.log_helper import my_logger

COUNT=2
class Demo:
    @my_logger(count=COUNT)
    def main(self):
        return "in main function"


if __name__ == '__main__':
    d = Demo()
    d.main()

ok,这样确实可以,我们还可以使用再简化一步

from functools import partial
from helper.log_helper import my_logger

COUNT=2
my_logger = partial(my_logger,count=2)


class Demo:
    @my_logger()
    def main(self):
        return "in main function"


if __name__ == '__main__':
    d = Demo()
    d.main()

暂时来看我们搞定了传参数的问题,这时候我们想如果外界调用了Demo类的main方法,并且向指定count的值怎么办呢?

我们知道外界调用Demo类传参的唯一途径就是向__init__里进行传参数,按照这个思路我们只能这么写了,

class Demo:
    def __init__(self):
        count =2
    @my_logger(count=self.count)
    def main(self):
        return "in main function"

但是这样并不可以,我们得到错误信息

NameError: name 'self' is not defined

在装饰器中无法使用self.形式的参数,难道这个问题解决不了么?

问题解决

在Python3.7之前确实没什么可行的方案。

我们知道在Python3.7的时候引入了dataclasses,我们可以通过它来简化__init__。

改下我们的代码

from functools import partial

from helper.log_helper import my_logger
from dataclasses import dataclass

@dataclass()
class Demo:
    count: int = 2
    logger: my_logger = partial(my_logger, count)

    @logger()
    def main(self):
        return "in main function"


if __name__ == '__main__':
    d = Demo()
    d.main()

如果使用Python3.8那么可以直接忽略掉dataclass

class Demo:
    count: int = 2
    logger: my_logger = partial(my_logger, count)

    @logger()
    def main(self):
        return "in main function"

这样我们就成功的解决了这个问题,突然想起来之前遇到的这个难题,现在算是解决了,希望对你有帮助。

后记

根据 @青南 大佬的提示,在装饰器里加个self也可以解决的。

import traceback
from functools import wraps

from loguru import logger


def my_logger(foo):
    @wraps(foo)
    def step2(self,*args, **kwargs):
        try:
            result = foo(self,*args, **kwargs)
            logger.info(f"{result=},{self.count=}")
            return result
        except Exception:
            logger.exception(traceback.format_exc())

    return step2

调用方

from helper.log_helper import my_logger

class Demo:
    def __init__(self):
        self.count= 2

    @my_logger
    def main(self):
        return "in main function"