python @register_Python中的注册器模块

简介

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

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

要注册的模块

models/model.py:

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17class Model:

pass

@Registers.model.register

class Model1(Model):

pass

@Registers.model.register

class Model2(Model):

pass

@Registers.model.register

class Model3(Model):

pass

注册器 Register1

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

37class 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():

等价于：

1func = 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类保存了名称->模块的数据，且提供了方便的注册装饰器。

所有注册器 Registers1

2

3

4

5

6class 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

29MODEL_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

6from 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

2Registers.model._dict before: {}

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

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

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