如何用装饰器包装 func.__code__.co_filename？

Posted 2023-03-23

【中文标题】如何用装饰器包装 func.__code__.co_filename？

【英文标题】：How to get wrapped func.__code__.co_filename with decorator?

【发布时间】：2021-11-15 20:24:36

【问题描述】：

一个简单的文件夹结构如下图，test.py中的所有函数都有关键字装饰器。

 lib
  |--- keyword.py
 main
  |--- test.py

关键字.py

from functools import wraps
def keyword(name=None, tags=(), types=()):

    def _method_wrapper(func):

        @wraps(func)
        def _passargs(self, *args, **kwargs):  
            print(func.__code__.co_filename)  # --> '..\main\test.py'
            return func(self, *args, **kwargs)

        print(_passargs.__code__.co_filename)  # --> '..\lib\keyword.py'
                
        return _passargs

    return _method_wrapper 

注意：print(..) 只是一个示例，我需要 _passargs 和 func 拥有相同的代码对象，而不是打印另一个变量 :)

如您所见，_passargs 弄错了co_filename。

这是来自robotframeworkkeyword.py，我出于自己的目的对其进行了修改。但是我无法弄清楚如何使_passargs 和func 都具有正确的源文件位置，以便robot.libdoc 可以正确生成doc.libspec。

谁能帮忙？

期待

func.__code__.co_filename = '..\main\test.py'
_passargs.__code__.co_filename = '..\main\test.py'

Python 版本 = 3.8.10

【问题讨论】：

改用func.__wrapped__.__code__.co_filename。如果您不确定func 是否被装饰，您可以执行try: func = func.__wrapped except AttributeError: pass 之类的操作。

嗨@kaya3，很抱歉没有解释导致您的误解。我实际上需要_passargs 来返回正确的__code__ 属性，但不仅要打印它！ robot.libdoc 是一个在初始化类时会生成关键字信息的函数，它不会进入_passargs，而只会进入_method_wrapper。因此，我需要_passargs返回实际的__code__.co_filename。

【参考方案1】：

在这里回答我自己的问题，以帮助可能遇到同样问题的人。

原来的装饰器函数：

from functools import wraps
def keyword(name=None, tags=(), types=()):

    def _method_wrapper(func):

        @wraps(func)
        def _passargs(self, *args, **kwargs):  
            print(func.__code__.co_filename)  # --> '..\main\test.py'
            return func(self, *args, **kwargs)

        print(_passargs.__code__.co_filename)  # --> '..\lib\keyword.py'
                
        return _passargs

    return _method_wrapper 

大多数时候，我们只需要最后一级装饰器（本例中为return func(self, *args, **kwargs)）来返回包装函数的正确属性。

但是，在我的例子中，有一个函数将在初始化类期间使用这个 keyword 装饰器扫描所有函数。而且它不会进入def _passargs 级别，因为我们没有提供任何参数，只是试图获取包装后的 func (return _passargs) 相关属性。

要解决此问题，我们只能在 RUNTIME 覆盖 co_filename。

这是在装饰器中覆盖 __code__ 对象的解决方案。

from types import CodeType
from functools import wraps


def keyword(name=None, tags=(), types=()):

    def _method_wrapper(func):

        @wraps(func)
        def _passargs(self, *args, **kwargs):  
            print(func.__code__.co_filename)  # --> '..\main\test.py'
            return func(self, *args, **kwargs)

        fix_co_filename(_passargs, func.__code__.co_filename)
        print(_passargs.__code__.co_filename)  # --> '..\main\test.py'
                
        return _passargs

    return _method_wrapper 


def fix_co_filename(func, co_filename):
    fn_code = func.__code__
    func.__code__ = CodeType(
        fn_code.co_argcount,
        fn_code.co_posonlyargcount,
        fn_code.co_kwonlyargcount,
        fn_code.co_nlocals,
        fn_code.co_stacksize,
        fn_code.co_flags,
        fn_code.co_code,
        fn_code.co_consts,
        fn_code.co_names,
        fn_code.co_varnames,
        co_filename,
        fn_code.co_name,
        fn_code.co_firstlineno,
        fn_code.co_lnotab,
        fn_code.co_freevars,
        fn_code.co_cellvars)