Python - 装饰器实现函数/类的注册

回顾一下装饰器

Python的装饰器语法想必大家都不陌生,所谓装饰器其实就是一个“函数增强器”,可以理解为输入一个函数对象,输出一个函数对象的函数。如果你的项目中需要在不更改一个函数函数体的前提下给这个函数增加额外的功能,那么就可以使用装饰器来完成。

举个简单的例子,我们的项目中目前有这么几个函数:

def add(a : int, b : int):
    return a + b

def multiply(a : int, b : int):
    return a * b

def minus(a : int, b : int):
    return a - b

我们现在希望无论在什么地方调用它们,都在控制台用黄色字体打印被调用函数的函数名。那么通过装饰器就可以轻松完成,我们首先需要实现装饰器函数:

def printFuncName(func):
    def decorator(*args, **kwargs):
        print(f"\033[33mcall {func.__name__}\033[0m")
        return func(*args, **kwargs)
    return decorator

然后我们将这个函数套在之前的三个函数上,以装饰器的形式“增强”这三个函数,并调用:

@printFuncName
def add(a : int, b : int):
    return a + b

@printFuncName
def multiply(a : int, b : int):
    return a * b

@printFuncName
def minus(a : int, b : int):
    return a - b

print(add(1, 2))
print(multiply(1, 2))
print(minus(1, 2))

out:

call add
3
call multiply
2
call minus
-1
由于知乎的文本框不是富文本编辑器,故原本五颜六色的输出这边都是一个颜色。=_=

而接下来我们就需要利用Python的装饰器语法完成函数/类的注册。


完成函数/类的注册

在Python中什么是函数和类的注册呢?这个问题我没有找到特别术语的解释,我结合个人的项目经历通俗讲讲我对注册的理解吧。如果你的项目中有许多可能会随处用到的函数和类,你希望把所有的这些指定的函数和类整合到一个字典(字典的key是函数名,类名或者自定义的名字,value是对应的函数对象或者类)中,那么这样一个过程称为注册。如果我将一个函数放入了这个字典中,那么就称完成了对该函数的注册。

这玩意儿的出现不是什么好玩的语法,而是利于真实开发而被人们使用的。那么什么情景下我们会需要使用这种注册机制呢?本蒟蒻目前大二,没做过什么大项目,只是以一点非常微弱的项目经历来谈谈个人的看法,我认为如果项目中出现了成批量的功能平行的函数或类时,使用注册机制能够提高项目的模块化。举两个例子:

  • 如果你在开发一款基于静态图的深度学习框架,那么肯定会实现许多的梯度函数,很显然,所有的梯度函数之间是功能平行的。那么在写计算图的反向传播时,我们就可以将注册好所有梯度函数的字典作为一个媒介,得到当前的运算节点的算子类型后,将这个类型的字符串作为key就可以获取对应的梯度函数了。比如我们目前反向传播的节点的算子类型为consumer_type="add",那么通过register[consumer_type]就可以得到add的梯度函数了。如果我们不适用注册技术,那么得到"add"后还得对梯度函数所在的python文件使用dir方法遍历每个attribute的__name__才能获取。而且,梯度函数存在上百个,全部存入一个字典中(也就是注册)也方便调用和管理。
  • 如果你在开发一款游戏,游戏中有上百张地图,每张地图都封装在一个class中,那么将所有的地图类进行注册后,在应用端接口,开发人员只需要获得注册字典,就可以方便地调取和管理所有的地图了。(最近在看gym的源代码,gym在封装强化学习环境时就是使用注册技术整合了所有的环境类,这里强化学习模拟器环境和3D游戏开发的地图如出一辙)

那么如何实现注册呢?

注册最简单的方式就是直接按照上面说的做:

register = {}

def func1():
    pass

f = lambda x : x

class cls1(object):
    pass

register[func1.__name__] = func1
register[f.__name__] = f
register[cls1.__name__] = cls1

但是这么做非常不优雅。而且对于较大的项目来说,破坏了项目的模块化,比如我们开发一个基于静态图的深度学习框架,那么你得在集中的100多行写上register[]=,非常麻烦,而且如果你需要删除函数或者修改某个函数的函数名时,连带着也需要修改register[]=。

那么有没有什么比较自动的方法呢?这就需要牵扯到之前讲的装饰器了,既然装饰器可以增强我们的函数,那么直接在装饰器中完成对传入的函数或者类的注册不就完事儿了吗?

我们接下来就是一个类Register,这个类的register方法充当我们的要实现注册功能的装饰器函数,由于注册的载体往往是字典,因此,我们这个Register类继承dict。说罢,咱们先简单实现写一下这个类的成员函数:

class Register(dict):
    def __init__(self, *args, **kwargs):
        super(Register, self).__init__(*args, **kwargs)
        self._dict = {}
    
    def register(self, target):
        pass
    
    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()

magic function和原本dict的几个方法都特别好写,关键就在于这个register函数,我们需要实现的装饰器函数功能为:如果用户不填入参数,那么就以被注册函数的函数名作为注册名,若填入参数,则以该参数作为该函数的注册名。

思路也很简单,我们只需要判断输入的参数是否可调用就可以判断用户是否往装饰器函数中输入了参数,然后做相应的处理就行,需要记住的是,无论是哪种情况,装饰器函数最终返回的都是函数对象:

    def register(self, target):
        def add_register_item(key, value):
            if not callable(value):
                raise Exception(f"register object must be callable! But receice:{value} is not callable!")
            if key in self._dict:
                print(f"warning: \033[33m{value.__name__} has been registered before, so we will overriden it\033[0m")
            self[key] = value
            return value

        if callable(target):            # 如果传入的目标可调用,说明之前没有给出注册名字,我们就以传入的函数或者类的名字作为注册名
            return add_register_item(target.__name__, target)
        else:                           # 如果不可调用,说明额外说明了注册的可调用对象的名字
            return lambda x : add_register_item(target, x)

fine,赶紧将register实例化吧!它现在就是一个强化版的注册字典。我就可以将实例后的对象的register方法作为装饰器函数套在我们想要注册的函数或者类上面啦!

register_functions = Register()

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

@register_functions.register("my multiply")
def multiply(a : int, b : int):
    return a * b

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

最后可以打印咱们的register看看:

if __name__ == "__main__":
    for k, v in register_functions.items():
        print(f"key: {k}, value: {v}")

out:

key: add, value: <function add at 0x000001FEBBD46C18>
key: my multiply, value: <function multiply at 0x000001FEBBD9B318>
key: minus, value: <function minus at 0x000001FEBBD96798>

为了少打几个字,我们可以再定义一下魔法方法__call__,最终的Register类如下:

class Register(dict):
    def __init__(self, *args, **kwargs):
        super(Register, self).__init__(*args, **kwargs)
        self._dict = {}
    
    def register(self, target):
        def add_register_item(key, value):
            if not callable(value):
                raise Exception(f"register object must be callable! But receice:{value} is not callable!")
            if key in self._dict:
                print(f"warning: \033[33m{value.__name__} has been registered before, so we will overriden it\033[0m")
            self[key] = value
            return value

        if callable(target):            # 如果传入的目标可调用,说明之前没有给出注册名字,我们就以传入的函数或者类的名字作为注册名
            return add_register_item(target.__name__, target)
        else:                           # 如果不可调用,说明额外说明了注册的可调用对象的名字
            return lambda x : add_register_item(target, x)
    
    def __call__(self, target):
        return self.register(target)
    
    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_functions = Register()

@register_functions
def add(a : int, b : int):
    return a + b

@register_functions("my multiply")
def multiply(a : int, b : int):
    return a * b

@register_functions
def minus(a : int, b : int):
    return a - b

print(register_functions)

out:

{'add': <function add at 0x000002707EB48948>, 'my multiply': <function multiply at 0x000002707EBC99D8>, 'minus': <function minus at 0x000002707EBCEA68>}

原文链接:Python进阶笔记(一)装饰器实现函数/类的注册 - 知乎

  • 30
    点赞
  • 21
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值