Python中的注册器机制

本文将会介绍Python中的注册器机制。

介绍

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

Python注册器机制本质上是用装饰器来实现的。

一个简单的例子

创建一个注册器（register.py）：

# -*- coding: utf-8 -*-
# @file: register.py
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()

对函数使用注册器：

from register import Register

my_register = Register()


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


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


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


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


@my_register
def custom_op(a, b):
    return a * b + a + b


if __name__ == '__main__':
    # check register information
    for k, v in my_register.items():
        print(f"key: {k}, value: {v}")
    print()

    # math open
    num1, num2 = 1, 2
    for func_name, func in my_register.items():
        print(f"{func_name} for {num1} and {num2}: ", func(num1, num2))

输出：

key: add, value: <function add at 0x104777d90>
key: multiply, value: <function multiply at 0x104939d80>
key: matrix multiply, value: <function mul_multiply at 0x10493a440>
key: minus, value: <function minus at 0x10493a3b0>
key: custom_op, value: <function custom_op at 0x104939cf0>

add for 1 and 2:  3
multiply for 1 and 2:  2
matrix multiply for 1 and 2:  None
minus for 1 and 2:  -1
custom_op for 1 and 2:  5

对类使用注册器

上面的简单例子是对函数实现了注册器机制，接下来我们看下如何对Python中的类实现注册器机制。

对类实现注册器：

# -*- coding: utf-8 -*-
# @file: class_register.py


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

    def __setitem__(self, key, value):
        if not callable(value):
            raise Exception(f"Value of a Registry must be a callable!\nValue: {value}")
        if key is None:
            key = value.__name__
        if key in self._dict:
            print("Key %s already in registry %s." % (key, self._name))
        self._dict[key] = value

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

    def register(self, target):
        """Decorator to register a function or class."""
        def add(key, value):
            self[key] = value
            return value

        if callable(target):
            # @reg.register
            return add(None, target)
        # @reg.register('alias')
        return lambda x: add(target, x)

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

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

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

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

实际中对类使用注册器：

from class_register import Register


class Registers:
    model = Register('model')
    dataset = Register('dataset')


class Model:
    pass


@Registers.model
class Model1(Model):
    pass


@Registers.model
class Model2(Model):
    pass


@Registers.model
class Model3(Model):
    pass


@Registers.dataset
class Dataset1:
    pass


print(Registers.model.items())
print(Registers.dataset.items())

输出结果为：

dict_items([('Model1', <class '__main__.Model1'>), ('Model2', <class '__main__.Model2'>), ('Model3', <class '__main__.Model3'>)])
dict_items([('Dataset1', <class '__main__.Dataset1'>)])

在实际的项目开发中，有不少项目使用了注册器机制，比如MMEngine，就实现了注册器，很好地管理了代码模块，参考网址为：https://mmengine.readthedocs.io/zh-cn/v0.7.0/advanced_tutorials/registry.html 。