Python类中如何为DataFrame概率计算方法实现通用校验装饰器

实现方案 #

把两类校验逻辑封装为可配置的类方法装饰器，替换原有存在重复执行问题、逻辑耦合的probability_zero装饰器，实现逻辑复用：

• 装饰器自动获取实例绑定的DataFrame，同时兼容位置传参、关键字传参场景
• 支持自定义配置需要校验的「列参数名」与对应「值参数名」配对，后续新增条件概率等多参数方法时无需修改装饰器逻辑
• 校验逻辑直接判断列、值是否真实存在，不再通过「概率值为0」反向推导，避免逻辑误判
• 校验通过后仅执行一次概率计算逻辑，消除原装饰器重复调用方法的性能浪费

完整实现代码 #

from functools import wraps
import inspect
import pandas as pd

class Probabilities():
    def __init__(self, dataframe):
        self.dataframe = dataframe
    
    def __repr__(self):
        return repr(self.dataframe)
    
    def __str__(self):
        return f"{self.dataframe}"
    
    def __getitem__(self, column):
        return self.dataframe[column]

    # 通用校验装饰器：传入(列参数名, 值参数名)的配对元组即可完成对应校验
    @staticmethod
    def validate_params(*check_pairs):
        def decorator(func):
            @wraps(func)
            def wrapper(self, *args, **kwargs):
                # 自动绑定所有传入参数，兼容位置传参、关键字传参
                sig = inspect.signature(func)
                bound_args = sig.bind(self, *args, **kwargs)
                bound_args.apply_defaults()
                params = bound_args.arguments

                # 逐组完成校验
                for col_arg, val_arg in check_pairs:
                    col = params.get(col_arg)
                    val = params.get(val_arg)
                    
                    # 校验1：列名是否存在于DataFrame
                    if col not in self.dataframe.columns:
                        print(f"Error: Column '{col}' not in table.")
                        return None
                    
                    # 校验2：目标值是否存在于对应列
                    if val not in self.dataframe[col].unique():
                        print(f"Error: Value '{val}' not in column '{col}'.")
                        return None
                
                # 所有校验通过后执行计算逻辑，仅执行一次
                return func(self, *args, **kwargs)
            return wrapper
        return decorator

    # 绝对概率计算：加装饰器指定要校验的column、value参数对
    @validate_params(("column", "value"))
    def _absolute_probability(self, column, value):
        bool_result = (self.dataframe[column] == value)
        return bool_result.mean()

    def absolute_probability(self, column, value):
        val = self._absolute_probability(column, value)
        if val is not None:
            formatted_val = f"""Probability of value: "{value}" in column: "{column}" is {val*100:.1f}%"""
            return formatted_val

    # 后续新增条件概率方法示例：直接加装饰器即可完成两组参数校验
    @validate_params(("target_col", "target_val"), ("condition_col", "condition_val"))
    def _conditional_probability(self, target_col, target_val, condition_col, condition_val):
        subset = self.dataframe[self.dataframe[condition_col] == condition_val]
        return (subset[target_col] == target_val).mean()
    
    def conditional_probability(self, target_col, target_val, condition_col, condition_val):
        val = self._conditional_probability(target_col, target_val, condition_col, condition_val)
        if val is not None:
            formatted_val = f"""Probability of value: "{target_val}" in column: "{target_col}" given "{condition_val}" in column: "{condition_col}" is {val*100:.1f}%"""
            return formatted_val

调用示例 #

# 测试数据
bank_df = pd.DataFrame({
    "jobb": ["management", "worker", "management", "teacher"]
})
p_df = Probabilities(bank_df)

# 正常调用
print(p_df.absolute_probability("jobb", "management"))
# 输出：Probability of value: "management" in column: "jobb" is 50.0%

# 列名不存在场景
print(p_df.absolute_probability("job", "management"))
# 输出：Error: Column 'job' not in table.
# 输出：None

# 值不存在场景
print(p_df.absolute_probability("jobb", "doctor"))
# 输出：Error: Value 'doctor' not in column 'jobb'.
# 输出：None

后续新增任何概率计算方法，只需要在内部计算函数上添加@validate_params()装饰器，传入需要校验的列、值参数名配对，即可自动复用两类校验逻辑，无需重复写校验代码。

内容的提问来源于stack exchange，提问作者matt.aurelio