如何在PySpark中编写自定义装饰器实现类方法的动态注册执行

• 你需要的装饰器要同时承担参数自动提取和方法增强的作用，另外可以用类扫描逻辑批量给所有带funct方法的类绑定装饰器，不需要手动修改原有类代码。
• 本方案默认你需要从DataFrame的首行提取funct所需的参数值，如果你是要做DataFrame转换逻辑，可以自行修改装饰器里的参数提取逻辑。

from pyspark.sql import SparkSession
import inspect

# ---------------------- 原有类（修正语法错误） ----------------------
class a:
    def __init__(self):
        pass
    def funct(self, a, b): # 原逻辑需要a、b两个参数
        return a + b

class b:
    def __init__(self):
        pass
    def funct(self, c, d): # 原逻辑需要c、d两个参数
        return c * d

# ---------------------- 自定义装饰器 ----------------------
# 装饰器带参数：传入参数与df列名的映射关系
def inject_df_params(col_mapping: dict):
    def decorator(func):
        def wrapper(self, df):
            # 从df首行提取对应列的值
            row = df.first()
            # 按照映射组装参数
            kwargs = {param: row[col_name] for param, col_name in col_mapping.items()}
            # 调用原funct方法
            return func(self, **kwargs)
        return wrapper
    return decorator

# ---------------------- 批量给目标类的funct方法加装饰器 ----------------------
# 配置每个类的参数与df列的映射
class_param_config = {
    "a": {"a": "col_a", "b": "col_b"}, # 类a的funct参数a对应df的col_a列，b对应col_b列
    "b": {"c": "col_c", "d": "col_d"}  # 类b的funct参数c对应df的col_c列，d对应col_d列
}

# 扫描当前模块所有类，自动绑定装饰器
for name, cls in inspect.getmembers(__import__(__name__), inspect.isclass):
    if name in class_param_config and hasattr(cls, 'funct'):
        # 替换原funct方法为加了装饰器的版本
        cls.funct = inject_df_params(class_param_config[name])(cls.funct)

# ---------------------- 动态调用示例 ----------------------
if __name__ == "__main__":
    # 初始化SparkSession
    spark = SparkSession.builder.appName("test_decorator").getOrCreate()
    # 测试用DataFrame
    test_data = [(1, 2, 3, 4)]
    test_df = spark.createDataFrame(test_data, schema=["col_a", "col_b", "col_c", "col_d"])
    
    # 动态实例化类并调用funct
    target_classes = [a, b]
    for cls in target_classes:
        instance = cls()
        result = instance.funct(test_df)
        print(f"类{cls.__name__}的funct执行结果：{result}")
    
    spark.stop()

• 类a的执行结果为1+2=3，类b的执行结果为3*4=12
• 如果需要扩展新的类，只需要在class_param_config里新增对应参数映射即可，不需要修改装饰器和调用逻辑

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