如何使用带有参数的python装饰器？

Question

我想定义一个装饰器，该装饰器将通过作为装饰器参数给出的名称注册类。我可以从 stackoverflow 和其他来源中阅读许多示例，这些示例展示了如何派生此类（棘手的）代码，但是当适应我的需要时，我的代码无法产生预期的结果。代码如下：

import functools

READERS = {}

def register(typ):
    def decorator_register(kls):
        @functools.wraps(kls)
        def wrapper_register(*args, **kwargs):
            READERS[typ] = kls 
        return wrapper_register
    return decorator_register

@register(".pdb")
class PDBReader:
    pass

@register(".gro")
class GromacsReader:
    pass

print(READERS)

此代码生成一个空字典，而我希望字典包含两个条目。你知道我的代码有什么问题吗？

Answer 1

接受参数（通过(...)）和装饰（通过@）都会导致函数调用。接受参数或装饰的每个“阶段”都映射到一个调用，从而映射到装饰器定义中的一个嵌套函数。 register 是一个三阶段的装饰器，它需要尽可能多的调用来触发其最里面的代码。其中，

• 第一个是参数（(".pdb")），
• 第二个是类定义 (@... class)，并且
• 第三个是类调用/实例化 (PDBReader(...))
  • 此阶段已中断，因为它没有实例化类。

为了将类本身存储在字典中，请将其存储在 second 阶段。由于不存储实例，因此删除第三阶段。

def register(typ):                # first stage: file extension
    """Create a decorator to register its target for the given `typ`"""
    def decorator_register(kls):  # second stage: Reader class
        """Decorator to register its target `kls` for the previously given `typ`"""
        READERS[typ] = kls 
        return kls                # <<< return class to preserve it
    return decorator_register

请注意，装饰器的结果会替换其目标。因此，您通常应该返回目标本身或等效对象。由于在这种情况下类会立即返回，因此无需使用functools.wraps。

READERS = {}

def register(typ):                # first stage: file extension
    """Create a decorator to register its target for the given `typ`"""
    def decorator_register(kls):  # second stage: Reader class
        """Decorator to register its target `kls` for the previously given `typ`"""
        READERS[typ] = kls 
        return kls                # <<< return class to preserve it
    return decorator_register

@register(".pdb")
class PDBReader:
    pass

@register(".gro")
class GromacsReader:
    pass

print(READERS)  # {'.pdb': <class '__main__.PDBReader'>, '.gro': <class '__main__.GromacsReader'>}

Answer 2

如果您实际上没有调用装饰器“包装”的代码，那么“内部”函数将不会触发，并且您不会在 READER 内创建条目。但是，即使您创建 PDBReader 或 GromacsReader 的实例，READER 内部的值也将属于类本身，而不是它们的实例。

如果你想做后者，你必须把wrapper_register改成这样：

def register(typ):
    def decorator_register(kls):
        @functools.wraps(kls)
        def wrapper_register(*args, **kwargs):
            READERS[typ] = kls(*args, **kwargs)
            return READERS[typ]
        return wrapper_register
    return decorator_register

我在类中添加了简单的 init/repr 以更好地可视化它：

@register(".pdb")
class PDBReader:
    def __init__(self, var):
        self.var = var
    def __repr__(self):
        return f"PDBReader({self.var})"

@register(".gro")
class GromacsReader:
    def __init__(self, var):
        self.var = var
    def __repr__(self):
        return f"GromacsReader({self.var})"

然后我们初始化一些对象：

x = PDBReader("Inside of PDB")
z = GromacsReader("Inside of Gromacs")
print(x) # Output: PDBReader(Inside of PDB)
print(z) # Output: GromacsReader(Inside of Gromacs)
print(READERS) # Output: {'.pdb': PDBReader(Inside of PDB), '.gro': GromacsReader(Inside of Gromacs)}

---

如果你不想将初始化的对象存储在READER中，你仍然需要返回一个初始化的对象，否则当你尝试初始化对象时，它会返回None。

然后您可以简单地将wrapper_register 更改为：

def wrapper_register(*args, **kwargs):
    READERS[typ] = kls
    return kls(*args, **kwargs)