如何用装饰器传递generate_sp500生成的sp500变量给sp500_rsi函数

核心方案：用装饰器自动注入sp500数据 #

装饰器的核心是提取start_date参数，调用generate_sp500生成数据，再将其传递给被装饰函数。下面是两种可靠的实现方式：

方式一：通过关键字参数传递（推荐） #

这种方式直观、易维护，被装饰函数只需声明sp500参数即可直接使用：

import pandas as pd

# 假设已实现的generate_sp500函数
def generate_sp500(start_date):
    # 替换为你的实际数据生成逻辑
    return pd.DataFrame({
        'date': pd.date_range(start_date, periods=10),
        'close': [150 + i*2 for i in range(10)]
    })

def sp500_decorator(func):
    def wrapper(*args, **kwargs):
        # 从参数中提取start_date（兼容位置/关键字参数）
        if 'start_date' in kwargs:
            start_date = kwargs['start_date']
        else:
            start_date = args[0]
        
        # 生成SP500数据
        sp500 = generate_sp500(start_date)
        
        # 将sp500作为关键字参数传入原函数
        return func(*args, **kwargs, sp500=sp500)
    return wrapper

# 改造后的sp500_rsi函数
@sp500_decorator
def sp500_rsi(start_date, rsi_period, sp500):
    # 直接使用sp500进行RSI计算
    delta = sp500['close'].diff()
    gain = delta.where(delta > 0, 0)
    loss = -delta.where(delta < 0, 0)
    avg_gain = gain.rolling(window=rsi_period).mean()
    avg_loss = loss.rolling(window=rsi_period).mean()
    rs = avg_gain / avg_loss
    rsi = 100 - (100 / (1 + rs))
    
    return rsi.dropna()

方式二：注入函数局部命名空间（无需修改参数列表） #

如果不想在原函数参数中添加sp500，可以用inspect模块修改函数局部作用域，但这种方式依赖Python内部机制，可读性稍差：

import pandas as pd
import inspect

def generate_sp500(start_date):
    return pd.DataFrame({
        'date': pd.date_range(start_date, periods=10),
        'close': [150 + i*2 for i in range(10)]
    })

def sp500_decorator(func):
    def wrapper(*args, **kwargs):
        # 解析函数签名，精准提取start_date
        sig = inspect.signature(func)
        bound_args = sig.bind(*args, **kwargs)
        bound_args.apply_defaults()
        start_date = bound_args.arguments['start_date']
        
        sp500 = generate_sp500(start_date)
        
        # 将sp500注入原函数的局部命名空间
        caller_frame = inspect.currentframe().f_back
        caller_frame.f_locals['sp500'] = sp500
        
        return func(*args, **kwargs)
    return wrapper

# 无需在参数中声明sp500
@sp500_decorator
def sp500_rsi(start_date, rsi_period):
    delta = sp500['close'].diff()
    # 后续RSI计算逻辑同上...
    return rsi.dropna()

常见问题排查 #

你之前的装饰器失效，大概率是以下原因：

• 没有正确从*args/**kwargs中提取start_date，导致generate_sp500调用失败
• 没有将生成的sp500注入到被装饰函数的作用域（既没传参也没修改局部命名空间）
• 参数传递顺序错误，导致原函数无法识别sp500变量

用第一种方式（关键字参数传递）能有效避免这些问题，同时保持代码的清晰性。

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