python - 如何从另一个函数中拦截一个函数中的变量

Question

假设我在 test.py 中有这样一个函数:

from math import sqrt

def a():
    intermediate_val = sqrt(4)
    return 5

def b():
    another_val = sqrt(9)
    return 8

我想编写一个函数来查看 a() 和 b() 并返回对 sqrt() 的任何调用的结果而不修改原始代码(装饰器会很好)。像这样:

import test

def intercept_value(fnc, intercept_fnc):
    # What goes here?

intercept_value('a', 'sqrt') == 2  # True
intercept_value('b', 'sqrt') == 3  # True

Answer 1

因为这既是一个有趣的问题又被标记为 abstract-syntax-tree，这里有一个解决方案，它将定位所有函数定义和在这些函数主体中调用的任何例程，然后计算子函数功能:

import ast, collections
class Parse:
    def __init__(self):
       self.f_subs = collections.defaultdict(dict)
    def memoize(self, sig):
       if all(isinstance(i, ast.Constant) for i in sig.args) and all(isinstance(i.value, ast.Constant) for i in sig.keywords):
           return {'args':[i.value for i in sig.args], 'keywords':{i.arg:i.value.value for i in sig.keywords}}
    def walk(self, tree, f = None):
       if isinstance(tree, ast.FunctionDef):
           f = tree.name
       elif isinstance(tree, ast.Call) and f is not None:
           if (m:=self.memoize(tree)) is not None:
              self.f_subs[f][tree.func.id] = m
       for i in getattr(tree, '_fields', []):
           v = getattr(tree, i)
           for j in ([v] if not isinstance(v, list) else v):
              self.walk(j, f)

import test
def intercept_value(fnc, intercept_fnc):
    p = Parse()
    with open(test.__file__) as f:
       p.walk(ast.parse(f.read()))
       return getattr(test, intercept_fnc)(*(f:=p.f_subs[fnc][intercept_fnc])['args'], **f['keywords'])
       
print(intercept_value('a', 'sqrt'))
print(intercept_value('b', 'sqrt'))

输出:

2
3