Python装饰器实现多样本量数据集模拟的疑问

问题背景 #

我正在尝试深入理解Python装饰器的工作原理，希望通过它实现函数多次运行，以模拟包含三个变量的数据集——这些数据集在样本量（20、50）及是否采用条件抽样上存在差异。

先构建用于抽样的总体分布：

from numpy.random import normal, negative_binomial, binomial
import pandas as pd

population_N = 100000
data = pd.DataFrame({ 
    "Variable A": normal(0, 1, population_N),
    "Variable B": negative_binomial(1, 0.5, population_N),
    "Variable C": binomial(1, 0.5, population_N)
})

原本通过重复代码生成不同样本：

sample_20 = data.sample(20)
sample_50 = data.sample(50)

condition = data["Variable B"] != 0

sample_20_non_random = data[condition].sample(20)
sample_50_non_random = data[condition].sample(50)

为简化流程，先编写了可指定抽样类型的函数：

def simple_function(data_frame, type = "random"):
    if (type == "random"):
        sample = data_frame.sample(sample_size)  # 原代码存在变量未定义问题
    else:
        condition = data_frame["Variable B"] != 0
        sample = data_frame[condition].sample(sample_size)
    return sample

尝试用装饰器替代循环实现多样本量支持，编写了如下代码，但无法理解装饰器的参数继承机制及迭代逻辑：

import functools

def decorator(cache = {}, **case):
    def inner(function):
        function_name = function.__name__
        if function_name not in cache:
            cache[function_name] = function
        @functools.wraps(function)
        def wrapped_function(**kwargs):
            if cache[function_name] != function:
                cache[function_name](**case)
            else:
                function(**case)
        return wrapped_function
    return inner

@decorator(sample_size = [20, 50])
def sample(data_frame, type = "random"):
    if (type == "random"):
        sample = data_frame.sample(sample_size)
    else:
        condition = data_frame["Variable B"] != 0
        sample = data_frame[condition].sample(sample_size)
    return sample

现有装饰器的核心问题 #

1. 变量作用域错误：被装饰的sample函数直接使用sample_size，但该变量既未在函数内部定义，也未从装饰器参数正确传递。
2. 逻辑偏离需求：cache的设计完全多余，既没有实现缓存功能，也未处理样本量迭代逻辑。
3. 缺失多样本遍历逻辑：装饰器接收了sample_size=[20,50]参数，但没有编写遍历列表、多次调用原函数的代码。

装饰器参数继承机制解析 #

带参数的装饰器本质是装饰器工厂，执行流程如下：

1. 当你写@decorator(sample_size=[20,50])时，Python会先执行decorator(sample_size=[20,50])，该函数接收参数后返回真正的装饰器函数（即代码中的inner）。
2. 接着Python将被装饰的sample函数传入inner，inner内部定义wrapped_function并返回。此时wrapped_function通过闭包特性，能访问到decorator传入的case参数（即sample_size=[20,50]）。
3. 之后调用sample(data, type="random")时，实际调用的是wrapped_function，它可以利用闭包中的参数控制原函数的执行逻辑。

实现多样本量迭代的正确装饰器 #

下面是满足需求的装饰器实现：它会遍历传入的样本量列表，对每个样本量调用原抽样函数，最终返回包含所有样本的字典。

import functools

def iterate_sample_sizes(**decorator_kwargs):
    # 提取装饰器传入的样本量列表
    sample_sizes = decorator_kwargs.get("sample_size", [])
    
    def inner_decorator(func):
        @functools.wraps(func)
        def wrapper(data_frame, type="random"):
            results = {}
            # 遍历每个样本量，调用原函数并保存结果
            for size in sample_sizes:
                sample_result = func(data_frame, type, size)
                # 生成清晰的结果键名
                key = f"{type}_{size}"
                results[key] = sample_result
            return results
        return wrapper
    return inner_decorator

# 修改原抽样函数，新增sample_size参数
@iterate_sample_sizes(sample_size=[20,50])
def sample(data_frame, type="random", sample_size=None):
    if type == "random":
        return data_frame.sample(sample_size)
    else:
        condition = data_frame["Variable B"] != 0
        return data_frame[condition].sample(sample_size)

使用示例 #

# 获取所有随机抽样样本
random_samples = sample(data, type="random")
# 获取所有条件抽样样本
non_random_samples = sample(data, type="non_random")

# 访问具体样本
print(random_samples["random_20"].shape)  # 输出(20, 3)
print(non_random_samples["non_random_50"].shape)  # 输出(50, 3)

内容的提问来源于stack exchange，提问作者Damon C. Roberts