每日一技: if-else 更优雅的替代方案

前言

Python 在 v3.10 之后的版本会引入 pattern matching 语法，这个语法会对消灭大量 if...elif...else 有巨大帮助。

但是在此之前，我们常使用的是 v3.9 之前的版本，并没有 pattern matching 语法。那么，在面对有大量 if..elif...else 的时候，有哪些方法可以让我们的代码更优雅丛容地处理这些场景呢？

在这里我提供两种解决方案，一种是字典，一种是装饰器。

字典

使用字典作为 if...else... 的替代方案是很多人的首选。

我们看看下面的例子：

def operations(operator, x, y):
    if operator == 'add':
        return x + y
    elif operator == 'sub':
        return x - y
    elif operator == 'mul':
        return x * y
    elif operator == 'div':
        return x / y

>>> operations('mul', 2, 8)
16

结合 lambda 和 dict 对上面的代码稍作修改，就能够得到下面这个更优雅的写法：

def operations(operator, x, y):
    return {
        'add': lambda: x+y,
        'sub': lambda: x-y,
        'mul': lambda: x*y,
        'div': lambda: x/y,
    }.get(operator, lambda: 'Not a valid operation')()

>>> operations('add', 1, 2)
3
>>> operations('sub', 1, 2)
-1 
>>> operations('unknown', 1, 2)
Not a valid operation

装饰器

from functools import wraps


def operator_dispatch(func):
    operators = dict()

    @wraps(func)
    def wrapper(arg0, *args, **kwargs):
        try:
            delegate = operators[arg0]
        except KeyError:
            pass
        else:
            return delegate(*args, **kwargs)
        
        return func(*args, **kwargs)
    
    def register(operator):
        def wrap(delegate):
            if operator in operators:
                raise ValueError('operator already registered')
            operators[operator] = delegate
            return delegate
        return wrap
    
    wrapper.register = register
    return wrapper

@operator_dispatch
def operator_handler(a, b):
    return f"You pass {a} and {b} but I don't know how to handle them."

@operator_handler.register("add")
def process_add(a, b):
    return a + b

@operator_handler.register("sub")
def process_sub(a, b):
    return a - b

@operator_handler.register("mul")
def process_mul(a, b):
    return a * b

@operator_handler.register("div")
def process_div(a, b):
    return a / b

if __name__ == "__main__":
    print(operator_handler("add", 1, 2))
    print(operator_handler("sub", 1, 2))
    print(operator_handler("mul", 1, 2))
    print(operator_handler("unknown", 1, 2))


"""Output
3
-1
2
You pass 1 and 2 but I don't know how to handle them.
"""

Gist

上述代码已经被我整理发布到 gist 中：