python - 在运行时修改函数代码

Question

我正在尝试编写一个装饰器，通过装饰器向函数添加详细日志记录(方法也不错，但我还没有尝试过)。这背后的动机是，在生产环境中将一行 add_logs 装饰器调用打补丁比添加 100 条调试行更容易(也更安全)。

例如:

def hey(name):
    print("Hi " + name)
    t = 1 + 1

    if t > 6:
        t = t + 1
        print("I was bigger")
    else:
        print("I was not.")
    print("t = ", t)
    return t

我想制作一个装饰器，将其转换为执行此操作的代码:

def hey(name):
    print("line 1")
    print("Hi " + name)

    print("line 2") 
    t = 1 + 1

    print("line 3")
    if t > 6:
        print("line 4")
        t = t + 1
        print("line 5")
        print("I was bigger")
    else:
        print("line 6")
        print("I was not.")
   print("line 7") 
   print("t = ", t)
   print("line 8")
   return t

到目前为止我得到了什么:

import inspect, ast
import itertools
import imp

def log_maker():
     line_num = 1
    while True:
        yield ast.parse('print("line {line_num}")'.format(line_num=line_num)).body[0]
        line_num = line_num + 1

def add_logs(function):
    def dummy_function(*args, **kwargs):
        pass
    lines = inspect.getsourcelines(function)
    code = "".join(lines[0][1:])
    ast_tree = ast.parse(code)
    body = ast_tree.body[0].body

    #I realize this doesn't do exactly what I want.
    #(It doesn't add debug lines inside the if statement)
    #Once I get it almost working, I will rewrite this
    #to use something like node visitors
    body = list(itertools.chain(*zip(log_maker(), body)))
    ast_tree.body[0].body = body
    fix_line_nums(ast_tree)
    code = compile(ast_tree,"<string>", mode='exec')

    dummy_function.__code__ = code
    return dummy_function

def fix_line_nums(node):
    if hasattr(node, "body"):
        for index, child in enumerate(node.body):
            if hasattr(child, "lineno"):
                if index == 0:
                    if hasattr(node, "lineno"):
                        child.lineno = node.lineno + 1
                    else:
                        # Hopefully this only happens if the parent is a module...
                        child.lineno = 1
                else:
                    child.lineno = node.body[index - 1].lineno + 1
            fix_line_nums(child)

@add_logs
def hey(name):
    print("Hi " + name)
    t = 1 + 1

    if t > 6:
        t = t + 1
        print("I was bigger")
    else:
        print("I was not.")
    print("t = ", t)
    return t

if __name__ == "__main__":
    print(hey("mark"))
    print(hey)

这会产生这个错误:

Traceback (most recent call last):
  File "so.py", line 76, in <module>
    print(hey("mark"))
TypeError: <module>() takes no arguments (1 given)

这是有道理的，因为编译创建了一个模块，当然模块不是可调用的。目前我已经尝试了一百种不同的方法来完成这项工作，但无法想出一个可行的解决方案。有什么建议吗？我是不是用错了方法？

(我还没有找到像这样在运行时实际修改代码的 ast 模块的教程。指向此类教程的指针也会有所帮助)

注意:我目前正在 CPython 3.2 上对此进行测试，但是 2.6/3.3_and_up 解决方案将不胜感激。目前，行为在 2.7 和 3.3 上是相同的。

Answer 1

当您编译源代码时，您会得到一个表示模块而不是函数的代码对象。将此代码对象替换为现有函数是行不通的，因为它不是函数代码对象，而是模块代码对象。不过，它仍然是一个代码对象，而不是一个真正的模块，因此您不能只执行 hey.hey 来从中获取函数。

相反，如 this answer 中所述，您需要使用exec 来执行模块的代码，将结果对象存储在字典中，并提取您想要的对象。您大致可以做的是:

code = compile(ast_tree,"<string>", mode='exec')
mod = {}
exec(code, mod)

现在 mod['hey'] 是修改后的函数。您可以将全局 hey 重新分配给它，或者替换它的代码对象。

我不确定你用 AST 做的事情是否完全正确，但无论如何你都需要做上面的事情，如果 AST 操作有问题，这样做会让你明白一点您可以从哪里开始调试它们。