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

Question

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

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

如果您实际上没有调用装饰器 "wrapping" 的代码，那么 "inner" 函数将不会触发，并且您不会在 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 中存储初始化对象，您仍然需要 return 一个初始化对象，否则当您尝试初始化对象时，它会 return None.

然后您只需将 wrapper_register 更改为：

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

Answer 2

接受参数（通过 (...)）和装饰（通过 @）都会导致函数调用。每个 "stage" 接受参数或修饰都映射到一次调用，因此在修饰器定义中映射到一个嵌套函数。 register 是一个三阶段装饰器，需要尽可能多的调用来触发其最内层代码。其中，

第一个是参数((".pdb")),
第二个是class定义(@... class)，
第三个是classcall/instantiation(PDBReader(...))
- 这个阶段被破坏了，因为它没有实例化 class。

为了在字典中存储class本身，存储在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

请注意，装饰器的结果会替换其目标。因此，您通常应该 return 目标本身或等效对象。由于在这种情况下 class 立即被 returned，因此无需使用 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'>}

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

how to use python decorator with argument?

python-3.x

python-decorators