Python 泛型函数 Singledispatch 使用及其扩充场景

@singedispatch

泛型函数可以通过传入参数类型不同而返回不同的处理结果，从而可以避免在代码中写入一系列的 if...elif...else 语句，这篇文章学习下 Python 中泛型函数的使用。

Python 中可以通过@functools.singledispatch将普通函数转为泛型函数 ^[1] ，📢 注意当使用@singledispatch定义函数时，注意 dispatch 是在第一个参数的类型上进行的。

简单定义函数示例如下

from functools import singledispatch
@singledispatch
def fun(arg, verbose=False):
    if verbose:
        print("Let me just say,", end=" ")
    print(arg)

然后向函数添加重载实现，请使用泛型函数的 register() 属性，该属性可以用作装饰器。对于带类型注释的函数，装饰器将自动推断第一个参数的类型:

@fun.register
def _(arg: int, verbose=False):
    if verbose:
        print("Strength in numbers, eh?", end=" ")
    print(arg)

@fun.register
def _(arg: list, verbose=False):
    if verbose:
        print("Enumerate this:")
    for i, elem in enumerate(arg):
        print(i, elem)

# 不添加类型形式
@fun.register(complex)
def _(arg, verbose=False):
    if verbose:
        print("Better than complicated.", end=" ")
    print(arg.real, arg.imag)

# 函数形式
def nothing(arg, verbose=False):
    print("Nothing.")
fun.register(type(None), nothing)

# 注册多种值
@fun.register(float)
@fun.register(Decimal)
def fun_num(arg, verbose=False):
    if verbose:
        print("Half of your number:", end=" ")
    print(arg / 2)

对应的输出结果如下：

>>> fun("Hello, world.")
Hello, world.
>>> fun("test.", verbose=True)
Let me just say, test.
>>> fun(42, verbose=True)
Strength in numbers, eh? 42
>>> fun(['spam', 'spam', 'eggs', 'spam'], verbose=True)
Enumerate this:
0 spam
1 spam
2 eggs
3 spam
>>> fun(None)
Nothing.
>>> fun(1.23)

从上面可以看出 singledispatch 可以从参数类型判断重载函数，满足 if type is int:... else ... 的功能，但明显的权限是从类型判断而不能从参数值判断，因此在某些场景下并不适用。

扩展形式

类似 singledispatch 的思路，主要需求是：可以根据参数值对函数进行重载，满足 if x==a: ... else ... 功能；^[2]

import functools


def valuedispatch(func):
    """Like functools.singledispatch() but dispatches by value of the first arg."""

    registry = {}

    @functools.wraps(func)
    def wrapper(arg0, *args, **kwargs):
        try:
            delegate = registry[arg0]
        except KeyError:
            pass
        else:
            return delegate(arg0, *args, **kwargs)

        return func(arg0, *args, **kwargs)

    def register(value):
        def wrap(func):
            if value in registry:
                raise ValueError(
                    f"@value_dispatch: existing handler " f"registered for {value!r}"
                )
            registry[value] = func
            return func

        return wrap

    def register_for_all(values):
        def wrap(func):
            for value in values:
                if value in registry:
                    raise ValueError(
                        f"@value_dispatch: existing handler " f"registered for {value!r}"
                    )
                registry[value] = func
            return func

        return wrap

    wrapper.register = register
    wrapper.register_for_all = register_for_all
    return wrapper

使用方法如下：

#!/usr/bin/env python3
# -*- encoding: utf-8 -*-
"""
@date: 2021-11-17 16:20:03
@author: dreamhomes.top
@description: test value dispatch.
"""
from dispatcher import valuedispatch


@valuedispatch
def eat(fruit):
    return f"I don't want a {fruit}..."


@eat.register("apple")
def _eat_apple(fruit):
    return "I love apples!"


@eat.register("eggplant")
@eat.register("squash")
# or @eat.register_for_all({'eggplant', 'squash'})
def _eat_what(fruit):
    return f"I didn't know {fruit} is a fruit!"

print(eat("apple"))  
>>> I love apples!
print(eat("squash"))
>>> I didn't know squash is a fruit!