Python中的注册器模块

Posted on Edited on In

简介

在一个稍大一点的python项目中，我们很有可能会用到注册器（register）。这个注册器不是用户账号注册的模块，而是项目中注册模块的一个模块。举个例子，一个深度学习项目可能支持多种模型；具体使用哪种模型可能是用户在配置文件中指定的。最简单的实现方式，就是维护一个模型名称->模型类的字典。但每当你增加一个模型时，这个字典就需要手动维护，比较繁琐。本文介绍一种注册器的模块，你需要维护的是需要注册的模块的代码路径（相对简介些）。

这个模块在我们的开源项目Delta中也有使用。

点这里看完整源代码

要注册的模块

models/model.py:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
class Model:
    pass


@Registers.model.register
class Model1(Model):
    pass


@Registers.model.register
class Model2(Model):
    pass


@Registers.model.register
class Model3(Model):
    pass

注册器 Register

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
class Register:

    def __init__(self, registry_name):
        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:
            logging.warning("Key %s already in registry %s." % (key, self._name))
        self._dict[key] = value

    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()

补充一个知识点，@是python的装饰器语法糖。

1
2
@decorate
def func():

等价于：

1
func = decorate(func)

这里，Register类似于一个dict（实际上是有一个_dict属性），可以set_item和get_item。关键是register函数，它可以作为装饰器，注册一个函数或者一个类。例如：

1
2
@register_obj.register
class Modle1:

等价于register_obj.register(Model1)，最终执行的是add(None, Model1)。

而：

1
2
@register_obj.register("model_one")
class Model1:

实际上是register_obj.register("model_one")(Model1)，最终执行的是add("model_one", Model_1)。

总结下：Register类保存了名称->模块的数据，且提供了方便的注册装饰器。

所有注册器 Registers

1
2
3
4
5
6
class Registers:

    def __init__(self):
        raise RuntimeError("Registries is not intended to be instantiated")

    model = Register('model')

Registers保存了所有的Register对象。

加载所有需要的模块到注册器 import_all_modules_for_register

在模块代码中加入注册装饰器之后，我们还需要把这些模块实际地导入，才能让这些子模块加入进注册器中。

一般大家会首先想到import。比如这里可以直接import models.models就可以让注册装饰器起作用。

但是import子模块这种形式很有可能导致循环引用的问题。为了避免循环引用，我们可以在代码入口处，统一地动态引入所有子模块。动态导入包使用importlib。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
MODEL_MODULES = ["models"]

ALL_MODULES = [("models", MODEL_MODULES)]


def _handle_errors(errors):
    """Log out and possibly reraise errors during import."""
    if not errors:
        return
    for name, err in errors:
        logging.warning("Module {} import failed: {}".format(name, err))


def import_all_modules_for_register(custom_module_paths=None):
    """Import all modules for register."""
    modules = []
    for base_dir, modules in ALL_MODULES:
        for name in modules:
            full_name = base_dir + "." + name
            modules.append(full_name)
    if isinstance(custom_module_paths, list):
        modules += custom_module_paths
    errors = []
    for module in modules:
        try:
            importlib.import_module(module)
        except ImportError as error:
            errors.append((module, error))
    _handle_errors(errors)

使用

最后我们使用下我们的注册器模块：

1
2
3
4
5
6
from register import import_all_modules_for_register
from register import Registers

print("Registers.model._dict before: ", Registers.model._dict)
import_all_modules_for_register()
print("Registers.model._dict after: ", Registers.model._dict)

输出：

1
2
Registers.model._dict before:  {}
Registers.model._dict after:  {'Model': <class 'models.models.Model'>, 'Model1': <class 'models.models.Model1'>, 'Model2': <class 'models.models.Model2'>, 'Model3': <class 'models.models.Model3'>}

可以看到，需要的模块已经加入到注册器中。

这个模块在我们的开源项目Delta中也有使用。

点这里看完整源代码