Python函数内调用拦截与修改：黑盒函数增强方案问询

假设有如下Python函数，func会调用add和mult：

def add(x, y):
    return x+y

def mult(x, y):
    return x*y

def func(x, y):
    added = add(x, y)
    multiplied = mult(x, y)
    return added, multiplied

需要实现类似装饰器的机制，拦截func内部对add和mult的调用：修改参数（比如固定y为10），并对结果做后处理（比如输出减半），让func等效于：

def modified_func(x, y):
    return (x+10)/2, (x*10)/2

核心约束 #

• 性能优先：涉及数百万次科学计算，不能有明显性能开销
• 通用黑盒方案：无需了解func内部结构（可能包含条件、循环，子函数调用顺序未知）
• 尽量避免全局修改add/mult
• 可对函数输入做预处理、输出做后处理

补充场景 #

实际待处理的函数结构如下（无法直接修改）：

def func(input):
    output = func1(input)
    output = func2(output)
    output = func3(output)
    # 更多函数调用
    output = funcn(output)
    return output

需要通过装饰器将其修改为：

def modified_func(input):
    output = func1(input, arg=some_arg)
    output = func2(output, arg=some_arg)
    output = func3(output, arg=some_arg)
    # 更多函数调用
    output = funcn(output, arg=some_arg)
    return output

1. 装饰器+局部函数替换（性能接近原生，推荐） #

利用Python的作用域覆盖特性，在装饰器内部重新定义被拦截的函数，让func在装饰后的局部作用域中调用这些替换后的函数，而非全局函数。这种方式性能几乎和原生调用一致，因为只是作用域查找的微小差异，没有额外运行时开销。

示例代码（针对最初的add/mult/func场景）：

def intercept_func(func):
    # 定义替换后的add：固定y=10，结果减半
    def add(x, y):
        return (x + 10) / 2
    # 定义替换后的mult：固定y=10，结果减半
    def mult(x, y):
        return (x * 10) / 2
    
    # 利用闭包让func在局部作用域中使用替换后的函数
    import types
    # 创建一个新的函数对象，替换其__globals__为包含替换函数的字典
    modified_func = types.FunctionType(
        func.__code__,
        {**func.__globals__, 'add': add, 'mult': mult},
        func.__name__,
        func.__defaults__,
        func.__closure__
    )
    return modified_func

# 使用装饰器
@intercept_func
def func(x, y):
    added = add(x, y)
    multiplied = mult(x, y)
    return added, multiplied

# 测试：func(5, 20) 会返回 (5+10)/2=7.5, (5*10)/2=25
print(func(5, 20))  # 输出 (7.5, 25.0)

针对补充场景（动态添加参数+预处理/后处理），可以扩展这个思路：

def inject_arg_and_wrap(target_funcs, arg_value, pre_process=None, post_process=None):
    def decorator(func):
        # 生成替换后的函数字典
        replaced_funcs = {}
        for func_name in target_funcs:
            original_func = func.__globals__[func_name]
            def wrapped(*args, original=original_func, arg=arg_value, pre=pre_process, post=post_process):
                # 预处理输入
                processed_args = args
                if pre:
                    processed_args = pre(*args)
                # 调用原函数，追加参数
                result = original(*processed_args, arg=arg)
                # 后处理输出
                if post:
                    result = post(result)
                return result
            replaced_funcs[func_name] = wrapped
        
        # 替换func的全局作用域
        modified_func = types.FunctionType(
            func.__code__,
            {**func.__globals__, **replaced_funcs},
            func.__name__,
            func.__defaults__,
            func.__closure__
        )
        return modified_func
    return decorator

# 使用示例：给func1-funcn都添加arg=some_arg，同时预处理输入、后处理输出
@inject_arg_and_wrap(
    target_funcs=['func1', 'func2', 'func3', 'funcn'],
    arg_value='some_arg',
    pre_process=lambda x: (x * 2,),  # 预处理：把输入x乘2
    post_process=lambda res: res + 1  # 后处理：结果加1
)
def func(input):
    output = func1(input)
    output = func2(output)
    output = func3(output)
    output = funcn(output)
    return output

优势：

• 性能极高：只是替换了函数的作用域，调用逻辑和原生一致，没有额外的装饰器嵌套开销
• 无需修改原函数代码，完全黑盒友好
• 不会全局修改原函数，避免影响其他代码

2. 字节码修改（性能最优，适合极端场景） #

如果追求极致性能，可以直接修改func的字节码，把内部的函数调用指令替换为我们的包装函数调用，或者直接修改参数和结果处理逻辑。这种方式完全没有运行时额外开销，但实现复杂度较高，需要对Python字节码有了解。

可以用bytecode库（安装命令：pip install bytecode）来简化操作，示例思路：

from bytecode import Bytecode, Instr
import types

def modify_bytecode(func):
    bc = Bytecode.from_code(func.__code__)
    # 遍历字节码，找到调用add/mult的指令（LOAD_NAME + CALL_FUNCTION）
    for i, instr in enumerate(bc):
        if instr.name == 'LOAD_NAME' and instr.arg in ('add', 'mult'):
            # 替换LOAD_NAME为我们的包装函数
            bc[i] = Instr('LOAD_NAME', f'wrapped_{instr.arg}')
    
    # 定义包装函数，加入到func的作用域
    def wrapped_add(x, y):
        return (x + 10)/2
    def wrapped_mult(x, y):
        return (x * 10)/2
    
    modified_code = bc.to_code()
    modified_func = types.FunctionType(
        modified_code,
        {**func.__globals__, 'wrapped_add': wrapped_add, 'wrapped_mult': wrapped_mult},
        func.__name__,
        func.__defaults__,
        func.__closure__
    )
    return modified_func

优势：

• 性能几乎和手写的modified_func完全一致，没有任何运行时额外开销
• 完全不影响全局函数

劣势：

• 实现复杂，需要理解Python字节码的栈模型
• 对不同Python版本的字节码兼容性可能有问题

3. 全局函数装饰（不推荐，除非无其他选择） #

如果可以接受全局修改的风险，也可以直接给add/mult加装饰器，但这会影响所有调用这些函数的代码，不符合优先需求：

def wrap_func(func_name, fix_y=10, post_process=lambda x: x/2):
    def decorator(func):
        def wrapped(x, y):
            # 固定y为10
            result = func(x, fix_y)
            # 后处理
            return post_process(result)
        return wrapped
    return decorator

# 全局修改add和mult
add = wrap_func('add')(add)
mult = wrap_func('mult')(mult)

# 此时调用func会自动生效，但其他地方调用add/mult也会被修改

• 优先选装饰器+局部函数替换：平衡了实现复杂度和性能，完全满足黑盒、无全局修改的需求
• 若追求极致性能，且能维护字节码逻辑，选字节码修改
• 全局装饰仅作为最后备选

内容的提问来源于stack exchange，提问作者Zet Khrush