仅使用NumPy实现带概率自动补全的trigger_randomly装饰器

import numpy as np
from functools import wraps

# 全局变量存储被装饰函数、对应显式概率、权重缓存
registered_funcs = []
explicit_probs = []
_cached_weights = None
_cached_func_count = -1

def trigger_randomly(probability=None):
    def decorator(fn):
        # 注册函数和对应概率，无参场景存None标记
        registered_funcs.append(fn)
        explicit_probs.append(probability)
        
        @wraps(fn)
        def wrapper(*args, **kwargs):
            global _cached_weights, _cached_func_count
            # 检测到新增被装饰函数时重算权重
            if _cached_func_count != len(registered_funcs):
                sum_explicit = sum(p for p in explicit_probs if p is not None)
                no_param_count = sum(1 for p in explicit_probs if p is None)
                no_param_weight = (1 - sum_explicit) / no_param_count if no_param_count > 0 else 0
                # 生成完整权重列表
                _cached_weights = [
                    p if p is not None else no_param_weight
                    for p in explicit_probs
                ]
                _cached_func_count = len(registered_funcs)
            
            # 用NumPy实现带权重随机选择
            selected_idx = np.random.choice(
                len(registered_funcs),
                p=_cached_weights
            )
            selected_fn = registered_funcs[selected_idx]
            return selected_fn(*args, **kwargs)
        return wrapper
    return decorator

# 测试用装饰器实例
@trigger_randomly(0.5)
def first():
    return "good person"

@trigger_randomly(0.3)
def second():
    return "kind person"

@trigger_randomly(0.1)
def third():
    return "happy person"

@trigger_randomly()
def fourth():
    return "beautiful person"

@trigger_randomly()
def fifth():
    return "reliable person"

# 分布统计测试
personality = {
    "good person": 0,
    "kind person": 0,
    "happy person": 0,
    "beautiful person": 0,
    "reliable person": 0
}

for i in range(1000):
    personality[fourth()] += 1

print(personality)

• 移除了random库的依赖，全程使用np.random.choice实现带权重的随机选择，符合仅用NumPy的要求。
• 取消无参装饰器的硬编码概率，动态计算权重：
  • 统计所有显式传入的概率总和
  • 计算剩余可分配概率：1 - 显式概率总和
  • 按无参装饰器的数量平分剩余概率，作为每个无参函数的权重
• 增加了权重缓存逻辑，只有当新增被装饰函数时才会重新计算权重，避免每次调用都重复计算，提升效率。

• 传入的显式概率总和不能超过1，否则会导致概率和不为1，触发NumPy的参数校验异常。
• 随机结果符合权重分布要求，多次调用的统计值会趋近于设定的权重比例。

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