从edgedb的源码学习如何简化大量if...elif...else

我们有一个变量s_fruit，需要根据其值来判断执行哪个逻辑，常规的写法就是if..elif...else的多重if语句，好一点就是把每个if给他提出去。

s_fruit = 'qqq'
if s_fruit == 'apple':
    print("I love apples!")
elif s_fruit in ['eggplant', 'squash']:
    print(f"I didn't know {s_fruit} is a fruit!")
else:
    print(f"I don't want a {s_fruit}...")

edge的源码里有一个装饰器，通过装饰器去分发给不同的分支。

源码如下：

import functools


def value_dispatch(func):

    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: there is already a 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: there is already a 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

使用装饰器来完成我们if语句的功能：

@value_dispatch
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')
def _eat_what(fruit):
    return f"I didn't know {fruit} is a fruit!"


print(eat('qqq'))

有其他的需求，可以自行改造源码