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("\n", "<br>\n")
    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("\n", "<br>\n")
    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.