Python如何根据变量值为装饰器动态传入不同参数

解决方案 #

有两种更简洁的规范实现方式，无需单独封装临时专用函数：

1. 内联条件表达式（单一场景首选） #

Python 允许装饰器位置直接使用合法表达式，直接写三元条件表达式即可满足需求：

@deviceCountAtLeast(1)
@(dtypes(torch.float) if NO_DOUBLE else dtypes(torch.float, torch.double))
def test_requires_grad_factory(self, devices, dtype):
    fns = [torch.ones_like, torch.testing.randn_like]
    x = torch.randn(2, 3, dtype=dtype, device=devices[0])

    for fn in fns:
        for requires_grad in [True, False]:
            output = fn(x, dtype=dtype, device=devices[0], requires_grad=requires_grad)
            self.assertEqual(requires_grad, output.requires_grad)
            self.assertIs(dtype, output.dtype)
            self.assertEqual(devices[0], str(x.device))

注：外层括号为可选，添加后可以提升条件判断的可读性

2. 通用条件装饰器（多重复用场景可选） #

如果项目中多处需要根据条件切换装饰器，可以封装通用工具，适用范围比临时专用函数更广：

def conditional_decorator(cond, decorator_true, decorator_false):
    return decorator_true if cond else decorator_false

# 使用逻辑和之前一致，函数名更通用适合多场景复用
@deviceCountAtLeast(1)
@conditional_decorator(NO_DOUBLE, dtypes(torch.float), dtypes(torch.float, torch.double))
def test_requires_grad_factory(self, devices, dtype):
    # 函数逻辑和原实现一致

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