如何对被function_2调用的external_lib模块中function_1进行Monkey Patch？

问题背景 #

有一个无法修改的外部库external_lib.py，其中function_2会调用同模块的function_1。需要在可修改的my_module.py中对function_1做猴子补丁，要求补丁后的函数保留原函数逻辑并新增自定义步骤，且让function_2调用时自动触发补丁后的函数。当前遇到两个问题：

• 如何对模块级函数进行猴子补丁；
• 如何让新函数在原函数基础上扩展逻辑。

原代码结构 #

external_lib.py（不可修改）

def function_1(arg1, arg2):
    # 耗时计算逻辑，根据arg1和arg2生成输出
    ...
    return output

def function_2(arg1, arg2):
    ...
    return function_1(arg1, arg2)

my_module.py（原始尝试）

import external_lib

def function_1(arg1, arg2):
    arg1 += '_'
    # 不知道如何复用原function_1的逻辑

# 尝试猴子补丁，但逻辑不完整
external_lib.function_1 = function_1

# 希望调用function_2时触发补丁后的function_1
external_lib.function_2('a', 'b')

解决方案 #

1. 模块级函数的猴子补丁方法

模块级函数的猴子补丁核心逻辑很直接：先保存原函数的引用，再将模块中的函数变量指向新的实现。你的原始尝试方向是对的，但缺少了保留原函数逻辑的关键步骤。

2. 扩展原函数逻辑的正确方式

要在原函数基础上新增步骤，必须先保留原函数的引用，然后在新函数中调用原函数，再插入自定义逻辑（可在原函数执行前、后或环绕执行）。

修正后的my_module.py代码 #

import external_lib

# 第一步：保存原函数的引用，防止补丁后丢失原有逻辑
original_func1 = external_lib.function_1

# 第二步：定义补丁函数，扩展原逻辑
def patched_func1(arg1, arg2):
    # 新增前置处理：给arg1加后缀
    arg1 += '_'
    # 调用原函数的核心逻辑，获取原始输出
    original_result = original_func1(arg1, arg2)
    # 新增后置处理：给原始输出加补丁标记
    patched_result = f"{original_result}_patched"
    return patched_result

# 第三步：执行猴子补丁，替换external_lib中的function_1
external_lib.function_1 = patched_func1

# 现在调用function_2会自动触发补丁后的function_1
final_result = external_lib.function_2('a', 'b')
print(final_result)

原理说明 #

external_lib.py中的function_2是直接引用模块全局作用域的function_1变量，当我们在my_module.py中修改external_lib.function_1的指向后，function_2调用的自然就是新的补丁函数，完全不需要修改外部库的代码。

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