【Python】Python的装饰器机制及注册器应用


当你的项目中需要成批量的函数和类,且这些函数和类功能上相似或并行时,为了方便管理,你可以把这些指定的函数和类整合到一个字典,你可以用函数名或类名作为字典的key,也可用使用自定义的名字作为 key,对应的函数或类作为 value。构建这样一个字典的过程就是注册,Python 引入注册器机制保证了这个字典可以自动维护,增加或删除新的函数或类时,不需要手动去修改字典。

Python 注册器机制本质上是用装饰器来实现的。下面我们将从基本的Python函数出发,逐步介绍装饰器,最后来学习注册器。


一、理解Python函数

首先定义一个函数,然后用不同的方式调用它

def math():
    return "I love math!"

print(math())
# output: 'I love math!'

# 这里math后面没有小括号,不是调用math函数
# 而是将math函数赋值给变量mathematics
mathematics = math

print(mathematics())
# output: 'I love math!'

在函数体中还可以定义函数,只是这个函数体内的函数不能在函数体外被直接调用:

def math():
    print('I am in the math() function!')

    def add():
        return 'I am in the add() function!'

    def multiply():
        return 'I am in the multiply() function!'

    print(add())
    print(multiply())
    print('I am now back to the math() function!')

math()
# output: 
# I am in the math() function!
# I am in the add() function!
# I am in the multiply() function!
# I am now back to the math() function!
add()
# NameError: name 'add' is not defined

函数体内的函数虽然不能在函数体外被直接调用,但是可以将它们返回出来。注意到返回的addmultiply后面没有小括号,那它就可以被传递,并且可以赋值给别的变量而被执行,如果有小括号,那它就会被执行。

def math(operation = 'add'):

    def add():
        return 'I am in the add() function!' 
    def multiply():
        return 'I am in the multiply() function!'

    if operation == 'add':
        return add
    else:
        return multiply

op = math()
print(op)
# output: <function math.<locals>.add at 0x7f0e64c9f560>

print(op())
# output: I am in the add() function!

我们还可以将函数作为参数传递给另一个函数:

def math(operation = 'add'):
    return 'I love math!'

def precompute(func):
    print('I am doing precomputations!')
    print(func())

precompute(math)

有了这样的印象之后,我们再写一个更加复杂一点的例子:

def a_decorator(func):

    def wrapTheFunction():
        print('Do some work before calling!')

        func()

        print('Do some work after calling')

    return wrapTheFunction

def math():
    print('I love math!')

math()
# outputs: I love math!

math = a_decorator(math)
math()
# output:
# Do some work before calling!
# I love math!
# Do some work after calling

二、理解Python装饰器

上一节的最后一个例子我们封装了一个函数,并且用另一个函数去修改这个函数的行为,这个功能其实就是Python装饰器所做的事情,只是我们以函数的形式显式的写了出来。Python中的装饰器提供了更简洁的方式来实现同样的功能,装饰器的写法是在被装饰的函数前使用@+装饰器名。现在我们用装饰器的写法来实现同样的功能:

def a_decorator(func):

    def wrapTheFunction():
        print('Do some work before calling!')

        func()

        print('Do some work after calling')

    return wrapTheFunction

@a_decorator 
def math():
    print('I love math!')

math()
# output:
# Do some work before calling!
# I love math!
# Do some work after calling

print(math.__name__)
# output: wrapTheFunction

但是与此同时,我们也发现了一个问题,那就是当我们输出被装饰函数的名字时,它被warpTheFunction函数替代了。Python为了解决这个问题,提供了一个简单的函数functools.wraps

from functools import wraps

def a_decorator(func):

    @wraps(func)
    def wrapTheFunction():
        print('Do some work before calling!')

        func()

        print('Do some work after calling')

    return wrapTheFunction

@a_decorator 
def math():
    print('I love math!')

math()
# output:
# Do some work before calling!
# I love math!
# Do some work after calling

print(math.__name__)
# output: math

不仅仅只有函数可以构建装饰器,类也可以用于构建装饰器,在构建装饰器类时,需要将原本装饰器函数的部分实现于__call__函数中即可:

from functools import wraps
class a_decorator(object):
    def __init__(self, func):
        self.func = func

    def __call__(self, *args, **kwargs):
        @wraps(self.func)
        def wrapTheFunction():
            print('Do some work before calling!')
            return self.func(*args, **kwargs)
        return wrapTheFunction

@a_decorator 
def math(*args, **kwargs):
    return 'I love math!'

m = math()
print(m())
# output:
# Do some work before calling!
# I love math!

三、学习Python注册器Registry

有了装饰器的基础之后,我们现在要走入注册器的世界了。Python的注册器本质上就是用装饰器的原理实现的。Registry提供了字符串到函数或类的映射,这个映射会被整合到一个字典中,开发者只要输入输入相应的字符串(为函数或类起的名字)和参数,就能获得一个函数或初始化好的类。 为了说明Registry的好处,我们首先看一下用一个字典存放字符串到函数的映射:

register = {}

def func():
    pass

f = lambda x : x

class cls(object):
    pass

register[func.__name__] = func
register[f.__name__] = f
register[cls.__name__] = cls

print(register)
# output:
# {'func': <function func at 0x7fec95a1add0>, '<lambda>': <function <lambda> at 0x7fec95a27a70>, 'cls': <class '__main__.cls'>}

这样做的缺点是我们需要手动维护register这个字典,当增加新的函数或类,或者删除某些函数或类时,我们也要手动修改register这个字典,因此我们需要一个可以自动维护的字典,在我们定义一个函数或类的时候就自动把它整合到字典中。为了达到这一目的,这里就使用到了装饰器,在装饰器中将我们新定义的函数或类存放的字典中,这个过程我们称之为注册。

这里我们需要定义一个装饰器类Register,其中核心部分就是成员函数register,它作为一个装饰器函数:

class Register(dict):
    def __init__(self, *args, **kwargs):
        super(Register, self).__init__(*args, **kwargs)
        self._dict = {}

    def register(self, target):
        def add_item(key, value):
            if not callable(value):
                raise Exception(f"Error:{value} must be callable!")
            if key in self._dict:
                print(f"\033[31mWarning:\033[0m {value.__name__} already exists and will be overwritten!")
            self[key] = value
            return value

        if callable(target):    # 传入的target可调用 --> 没有给注册名 --> 传入的函数名或类名作为注册名
            return add_item(target.__name__, target)
        else:                   # 不可调用 --> 传入了注册名 --> 作为可调用对象的注册名 
            return lambda x : add_item(target, x)

    def __setitem__(self, key, value):
        self._dict[key] = value

    def __getitem__(self, key):
        return self._dict[key]

    def __contains__(self, key):
        return key in self._dict

    def __str__(self):
        return str(self._dict)

    def keys(self):
        return self._dict.keys()

    def values(self):
        return self._dict.values()

    def items(self):
        return self._dict.items()

将Register实例化,然后打印验证一下:

register_func = Register()

@register_func.register
def add(a, b):
    return a + b

@register_func.register
def multiply(a, b):
    return a * b

@register_func.register('matrix multiply')
def multiply(a, b):
    pass

@register_func.register
def minus(a, b):
    return a - b

@register_func.register
def minus(a, b):
    return a - b

for k, v in register_func.items():
    print(f"key: {k}, value: {v}")
# output:
# Warning: minus already exists and will be overwritten!
# key: add, value: <function add at 0x7fd18ee7cb90>
# key: multiply, value: <function multiply at 0x7fd18ee95170>
# key: matrix multiply, value: <function multiply at 0x7fd18ee95320>
# key: minus, value: <function minus at 0x7fd18ee95200>

如果不想手动调用register()函数,可以在Register类中添加一个__call__()函数:

class Register(dict):
    def __init__(self, *args, **kwargs):
        super(Register, self).__init__(*args, **kwargs)
        self._dict = {}

    def __call__(self, target):
        return self.register(target)

    def register(self, target):
        def add_item(key, value):
            if not callable(value):
                raise Exception(f"Error:{value} must be callable!")
            if key in self._dict:
                print(f"\033[31mWarning:\033[0m {value.__name__} already exists and will be overwritten!")
            self[key] = value
            return value

        if callable(target):    # 传入的target可调用 --> 没有给注册名 --> 传入的函数名或类名作为注册名
            return add_item(target.__name__, target)
        else:                   # 不可调用 --> 传入了注册名 --> 作为可调用对象的注册名 
            return lambda x : add_item(target, x)

    def __setitem__(self, key, value):
        self._dict[key] = value

    def __getitem__(self, key):
        return self._dict[key]

    def __contains__(self, key):
        return key in self._dict

    def __str__(self):
        return str(self._dict)

    def keys(self):
        return self._dict.keys()

    def values(self):
        return self._dict.values()

    def items(self):
        return self._dict.items()


register_func = Register()

@register_func
def add(a, b):
    return a + b

@register_func
def multiply(a, b):
    return a * b

@register_func('matrix multiply')
def multiply(a, b):
    pass

@register_func
def minus(a, b):
    return a - b

for k, v in register_func.items():
    print(f"key: {k}, value: {v}")
# output:
# key: add, value: <function add at 0x7fdedd53cb90>
# key: multiply, value: <function multiply at 0x7fdedd540200>
# key: matrix multiply, value: <function multiply at 0x7fdedd5403b0>
# key: minus, value: <function minus at 0x7fdedd540320>