Python——MonkeyPatch使用笔记

---

整体说明

• 猴子补丁（Monkey Patch）可以灵活的修改已经定义的函数/类名，常用于对开源项目进行二次开发，示例如下

  import module
  module.func = new_func

• 但是 猴子补丁使用过程中，可能遇到很多比较隐晦的错误，本文记录一些事项

---

使用注意前置总结

• 在执行 module.func = new_func 赋值之前，建议全局搜索（grep）目标函数名，逐一确认以下风险点：
  • 检查所有调用该函数的地方，看看是否有文件在顶部使用了 from module import func 的写法
    • 如果存在这种写法，请追踪该文件的导入时机，确认它是在补丁代码之后被导入执行的，还是之前
      • 如果是之前，补丁会失效
  • 确认目标函数是否会在子进程或 Ray Worker 中被调用
    • 如果是，必须确认补丁逻辑已在那些子进程中重复执行
  • 检查新函数内部是否意外调用了被修补的模块属性
    • 如果是，务必使用“预先捕获原函数”的技巧避免无限递归

---

Monkey Patch 核心的本质

• 本质：module.func = new_func 替换的不是函数本身这个对象，而是模块对象 __dict__ 字典中名为 func 的那个条目
• 补丁是否生效，不取决于“在哪里”调用，完全取决于调用方以何种方式、在什么时机 去获取这个函数的名字

---

四种引用场景的生效性（特别注意）

场景一：被修补模块内部的裸名调用，或者调用方使用 import module 方式

• 这种情况下，每次调用 都会动态地去模块的 __dict__ 中查找 func 这个名字
  • 注意：是每次调用都动态查找，不是每次导入动态查找
    • 也就是说：导入 -> 执行补丁 -> 调用，补丁依然生效
• 修改了字典里的值，后续的调用自然拿到新函数，因此执行猴子补丁以后得所有调用均 生效 新函数
• 补充知识：
  • Python 模块的初始化不是从“被调用”开始的，而是从“被导入”开始的
  • 如果某个模块从未被任何地方导入，它就永远不会被初始化
• 特别注意：如果在执行补丁这句话前就使用 import module 来加载函数，那么这里（导入 -> 执行补丁 -> 调用）仍然会调用到新函数
  • 只要调用是发生在补丁之后就行（PS：补丁之前的调用是不会生效的！）

场景二：调用方在补丁操作 之后 执行 from module import func

• 导入语句会在执行的瞬间去模块的 __dict__ 中读取 func 的当前值
• 如果这行导入代码写在补丁赋值语句的后面，那么它读取到的就是已经挂载好的新函数，因此补丁 生效

场景三：调用方在补丁操作 之前 执行 from module import func

• 这是导致补丁失效的经典元凶
  • from ... import ... 的本质是“值拷贝”（即绑定对象引用）
  • 当导入执行时，Python 将原函数的内存地址直接复制给了当前模块的局部变量 func
  • 此后，即使源模块的 __dict__ 被修改，这个局部变量依然固执地指着旧函数的地址，因此补丁 失效

场景四：补丁后的代码对旧函数有回调需求

• 如果需要在新函数内部调用原始逻辑，千万不能用 module.func()（会无限递归）或 from module import func（会拿到新函数）
• 正确的做法是：在打补丁赋值之前 ，先用 from module import func 把原始函数捕获到一个独立的局部变量中，供新函数回调使用

---

最佳实践

延迟导入（Lazy Import），规避时机陷阱

• 如果模块在启动时就被加载，且执行了 from ... import ...，补丁将无法干预
• 解决办法是将导入语句移到函数内部 ，让导入动作推迟到函数实际被调用时才发生（即不在文件开始导入！）
• 只要补丁代码在函数调用之前被执行，延迟导入就能确保拿到修补后的版本

  def my_compute():
      from module import func  # 调用时才查表
      func(...)

预先捕获原函数，安全实现“套娃”式增强

• 当需要给一个函数增加新逻辑，又必须回调原始实现时，可以利用“补丁前导入生效”的特性
• 在补丁赋值的前一行，用 from ... import ... 将原函数存为别名，这样新函数内部可以安全回调，绝不会触发无限递归

  # patch_script.py（在赋值操作前被执行）
  from module import func as original_func
  
  def my_func(data, **kwargs):
      print("执行增强逻辑...")
      return original_func(data, **kwargs)  # 安全回调原函数
  
  # 执行补丁
  import module
  module.func = my_func

使用 functools.partial 注入配置参数

• 如果修补的新函数需要额外的配置参数，但调用方又没有传递该参数，强行修改调用方会破坏兼容性
• 此时可以用 partial 将配置固化绑定，生成一个签名与原函数完全一致的新可调用对象，调用方完全无感知

  from functools import partial
  
  def my_func_with_config(data, extra_cfg, **kwargs):
      print(f"额外配置: {extra_cfg}")
      # ...
  
  mod.func = partial(my_func_with_config, extra_cfg=global_config)
  # 调用方执行 mod.func(data, **kwargs) 时，extra_cfg 已自动注入

---

特别注意：分布式与多进程 Monkey Patch 不生效

• 核心原因：补丁是进程级别的，不能跨进程传播
  • 无论是使用 multiprocessing、torchrun 还是 Ray，当启动子进程或 Worker 进程时，操作系统会复制父进程的内存空间（或启动全新的解释器）
  • 如果在主进程中打了补丁，但目标函数在 Worker 进程中执行，该 Worker 并没有执行补丁代码，因此调用到的依然是原始函数
• 解决办法：
  • 必须确保补丁代码在每个 Worker 进程的初始化阶段都被执行
  • 具体可以通过三种途径实现：
    • 将补丁代码写在 Worker 类的 __init__ 方法中
    • 将补丁代码放在一个独立的 .py 文件中，在 Worker 启动时通过 runtime_env 或启动脚本强制导入
    • 如果是 Ray，可以利用 ray.remote 的 init_hook 参数或装饰器的初始化函数来执行补丁