Python中的注册器模块(转)

简介

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

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

点这里看完整源代码

要注册的模块

models/model.py:

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

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的装饰器语法糖。

@decorate def func():

等价于：

func = decorate(func)

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

@register_obj.register class Modle1:

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

而：

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

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

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

所有注册器 Registers

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。

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)

使用

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

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)

输出：

Registers.model._dict before: {} Registers.model._dict after: {'Model': , 'Model1': , 'Model2': , 'Model3': }

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

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