如何让装饰器生成可被pickle序列化的Python函数

问题现象 #

使用functools.wraps装饰函数后，若将装饰后的函数赋值回原函数名，pickle.dumps()可正常执行；但如果赋值给新变量，就会抛出PicklingError错误。

示例代码1（执行正常）：

import pickle
from functools import wraps

def wrapper(func):
    @wraps(func)
    def _func(*args, **kwargs):
        print(args, kwargs)
        return func(*args, **kwargs)
    return _func

def foo(a: int):
    return a + 1

foo = wrapper(foo)
pickle.dumps(foo)  # 无报错

示例代码2（抛出错误）：

def foo(a: int):
    return a + 1

_foo = wrapper(foo)
pickle.dumps(_foo)  # 抛出PicklingError

错误信息：

PicklingError: Can't pickle <function foo at 0x12e116b00>: it's not the same object as __main__.foo

问题原因 #

@wraps(func)会把原函数的__name__、__qualname__等元属性复制到装饰后的包装函数上。而pickle序列化函数时，会通过函数的__module__和__qualname__去模块中查找对应的函数对象：

• 当装饰后的函数赋值回原函数名foo时，模块中的__main__.foo就是这个装饰后的函数，两者匹配，序列化正常。
• 当赋值给新变量_foo时，装饰后的函数的__qualname__还是原函数的foo，但模块中的__main__.foo是未被装饰的原函数，两者不是同一个对象，因此pickle报错。

解决方案 #

修改装饰器，通过解析调用上下文的赋值语句，自动将包装函数的__name__和__qualname__设置为目标变量名，无需手动修改。

修改后的装饰器代码：

import pickle
from functools import wraps
import inspect

def wrapper(func):
    @wraps(func)
    def _func(*args, **kwargs):
        print(args, kwargs)
        return func(*args, **kwargs)
    
    # 获取调用者上下文，解析赋值目标变量名
    caller_frame = inspect.currentframe().f_back
    try:
        # 提取调用行的源代码
        code_line = inspect.getframeinfo(caller_frame).code_context[0].strip()
        # 解析赋值语句左侧的变量名
        if '=' in code_line:
            target_name = code_line.split('=')[0].strip()
            # 更新包装函数的名称属性
            _func.__name__ = target_name
            _func.__qualname__ = target_name
            # 重置__wrapped__属性，确保属性关联正确
            delattr(_func, '__wrapped__')
            setattr(_func, '__wrapped__', func)
    finally:
        # 删除帧引用，避免内存泄漏
        del caller_frame
    
    return _func

# 测试用例1：赋值回原函数名
def foo(a: int):
    return a + 1

foo = wrapper(foo)
pickle.dumps(foo)  # 正常执行

# 测试用例2：赋值给新变量
def bar(a: int):
    return a + 2

_bar = wrapper(bar)
pickle.dumps(_bar)  # 正常执行

说明 #

1. 借助inspect模块获取调用者的源代码行，解析出赋值目标变量名。
2. 更新包装函数的__name__和__qualname__为目标变量名，让pickle能找到正确的函数对象。
3. 重新设置__wrapped__属性，清除@wraps缓存的旧属性关联。
4. 手动删除帧引用，避免潜在的内存泄漏问题。

注意：该方法依赖赋值语句为单行格式，若遇到链式赋值、多行赋值等复杂场景，需调整解析逻辑，但能覆盖绝大多数常规使用场景。

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