Python中singledispatch装饰器实现函数重载

本文参照"流畅的Python"这本书有关于singledispatch实现函数重载的阐述^[1].

假设我们现在要实现一个函数, 功能是将一个对象转换成html格式的字符串. 怎么重载呢?

你可能会想, 用什么装饰器, Python不是动态类型么, 写成如下这样不就重载了嘛? 

def htmlize(obj):
    content = html.escape(repr(obj))
    return '<pre>{}</pre>'.format(content)

这个函数可以接受任意类型的参数. 你看,这不就重载了么?

如果我想让不同类型的对象有不同形式的html字符串呢 ? 你可能会说, 那就加个类型判断呗! 像下面这样.

def ifelseHtmlize(obj):
    if isinstance(obj, str):
        content = html.escape(repr(obj))
        content = '<pre>{}</pre>'.format(content)
    elif isinstance(obj, numbers.Integral):
        content = "<pre>{0} (0x{0:x})</pre>".format(obj)
    
    return content

额...这样当然可以实现上述需求. 然而, 当每种类型的处理逻辑比较复杂时, 以上方法大大增加了函数的篇幅, 影响可读性和可维护性. 你可能又会说, 那把每个分支抽象成函数就好了, 这样ifelseHtmlize函数的代码量就少了. 但是这样需要抽象出很多名字不一样的函数, 维护起来还是不太容易. 

因此, 我们需要寻求一种方法, 让每个重载函数能够关注自身需要处理的类型. 而且可以很简单的添加和去除. 于是就有了singledispatch这个装饰器. 先简单介绍下装饰器, 装饰器本质上是个函数, 其输入是一个函数, 返回值也是个函数. 把"@装饰器"放到哪个函数的头顶, 哪个函数就会作为参数输入到装饰器, 返回的函数再赋值给被装饰的函数. 这个技术的目的是对被装饰的函数做一些其他的手脚, 使其具有一些需要的特性. 例如singledispatch是Python内置众多装饰器之一, 就是为了达到函数重载的目的.

如何用singledispatch实现重载呢, 直接上代码:

@singledispatch
def htmlize(obj):
    content = html.escape(repr(obj))
    return '<pre>{}</pre>'.format(content)

@htmlize.register(str)
def _(text):
    content = html.escape(text).replace("
", "<br>
")
    return '<pre>{}</pre>'.format(content)

@htmlize.register(numbers.Integral)
def _(n):
    return "<pre>{0} (0x{0:x})</pre>".format(n)

乍一看脑壳疼. 简化一下, 如果去掉函数头顶上的装饰器, 再将这三个函数看作是一个函数名, 这和C++里的重载是类似的. 什么? 没学过C++, 那这句话当我没说, 接着看.

首先看第一个函数, 之前实现的htmlize函数被singledispatch装饰了, 意味着htmlize再也不是原来那个单纯的htmlize了. 现在的htmlize函数, 是被传入singledispatch后再返回出来的那个对象/函数. 具体的我们看一下singledispatch对htmlize都干了什么.

def singledispatch(func):
    # 用来记录 类型->函数
    registry = {}
    ...
    
    # 用来获得 指定类型的重载函数
    def dispatch(cls):
    ...
    
    # 注册新的 类型->函数
    def register(cls, func=None):
    ...
    
    # 重载函数入口
    def wrapper(*args, **kw)
    ...
    
    # 默认的重载函数
    registry[object] = func
    wrapper.register = register
    wrapper.dispatch = dispatch
    ...

    # 返回 wrapper供使用这调用重载函数
    return wrapper    

在singledispatch函数中, 参数func就是原始的htmlize函数. 然后先定义了一个字典registry, 看这个名字大概能猜出这是干嘛的. 对了, 就是注册用的, 注册啥呢? 当然是新的重载函数. 接着我们跳过中间几个函数, 先看下面一句registry[object] = func, 这将object类型对应的重载函数设置为func, 也就是传入的htmlize. 为什么是object呢? 这里埋个伏笔, 后面再讲. 现在说一下跳过的那三个函数, 实际上是三个闭包:

dispatch: 输入参数cls是类型, 根据给定的cls去registry中找对应的重载函数, 然后返回给调用者. 后面会见到这个函数是如何被调用的.

register: 用来注册新的重载函数. 核心的功能就是向registry中添加新的 类型->函数.

wrapper: 重载函数调用的入口, 负责执行正确的重载函数, 并返回结果.

最后对wrapper函数赋值了两个属性register和dispatch, 分别是同名的两个闭包, 接着返回wrapper成为新的htmlize. 什么, 没听说过直接给函数属性赋值? 看PEP 232 -- Function Attributes. 这里重点关注一下register, 有了它, 用户可以通过wrapper调用这个闭包. 意味着用户可以用register注册新的重载函数. 我们最终的目的要呼之欲出啦!

到这里只是对htmlize函数进行重载的使能. 接下来就可以定义htmlize的重载函数了. 直接上代码:

@htmlize.register(str)
def htmlize4str(text):
    content = html.escape(text).replace("
", "<br>
")
    return '<pre>{}</pre>'.format(content)

@htmlize.register(numbers.Integral)
def htmlize4Integral(n):
    return "<pre>{0} (0x{0:x})</pre>".format(n)

上面两个函数就是htmlize的两个重载函数, 分别用于处理字符串和Integral类型. 这两个函数均被htmlize.register(cls)返回的函数装饰. 我们知道, 上述singledispatch返回的wrapper会重新赋值给htmlize, 所以调用htmlize.register(cls)即是调用闭包register. 我们以htmlize.register(str)为例, 看看闭包register干了什么:

    def register(cls, func=None):
        ...
        if func is None:
            return lambda f: register(cls, f)
        registry[cls] = func
        ...
        return func

调用register(str), 因为func是None, 所以进入分支, 直接返回一个函数lambda f: register(str, f). 返回之后的函数作为装饰器, 对htmlize4str函数进行装饰. 于是htmlize4str函数作为lambda表达式中的f, 实际上调用了register(str, htmlize4str). 于是, 又回到了上述函数, 这次func==htmlize4str非None, 于是str->htmlize4str得以注册, 最后返回htmlize4str.

以上, 就完成了注册重载函数的过程了. 那如何实现传入htmlize不同参数, 执行不同的函数呢. 比如调用htmlize("23333"), 如何定位到htmlize4str("23333")呢? 现在回忆一下singledispatch装饰的htmlize现在是什么? 是wrapper闭包啊, 所以调用htmlize("23333"), 即调用wrapper("23333"). 我们看看wrapper做了什么:

    def wrapper(*args, **kw):
        return dispatch(args[0].__class__)(*args, **kw)

wrapper将输入的第一个参数的__class__, 即类型输入到dispatch. 我们之前提到过, dispatch这个函数用于找到指定类型的重载函数. dispatch返回后执行这个重载函数, 再将结果返回给调用者. 例如, wrapper("23333")首先调用dispatch(str), 因为"23333"的类型是str, 找到对应的重载函数, 即htmlize4str, 然后再调用htmlize4str("23333"). 实现重载啦啦啦! 而且我们发现重载函数的函数名, 对于调用htmlize是透明的, 根本用不到. 所以重载的函数名可以用_替代, 这样更好维护代码.

最后我们看看dispatch这个闭包干了些什么:

    def dispatch(cls):
        """generic_func.dispatch(cls) -> <function implementation>

        Runs the dispatch algorithm to return the best available implementation
        for the given *cls* registered on *generic_func*.

        """
        ...
            try:
                impl = registry[cls]
            except KeyError:
                impl = _find_impl(cls, registry)
        ...
        return impl

核心的功能就是到registry里找对应的cls类型的重载函数, 然后返回就行了. 那如果没找到呢? 比如我调用了htmlize({1, 2, 3}), 这时cls是list类型, 在registry里没有找到对应的重载函数咋办呢? 在上述代码中, 捕捉了registry抛出的KeyError异常, 即在没有找到时执行_find_impl(cls, registry), 这又是干嘛的呢? 这里不展开讲了, 我也展不开. 总之, 用一句话来说: 找到registry中和cls类型最匹配的类型, 然后返回其重载函数.

看看现在我们的registry里有哪些类型呢? str, numbers.Integral. 哦!!! 还有object (回忆一下, 在singledispatch中有这么一句: registry[object] = func). Python里所有类型都继承object类型, 于是返回registry中object对应的重载函数, 即最原始的htmlize. 

以上.

[1] Ramalho, Luciano. Fluent Python : clear, concise, and effective programming. Sebastopol, CA : O'Reilly, 2015.