如何简化大量的 if…elif…else 代码？

Posted 2021-12-11 Python猫

花下猫语：在日常代码中，我们总是会面对大量 if...elif...else 条件分支选择的问题。说实话，我在多数情况下首推就是老老实实地写 if...elif，并且设法把每个分支下的内容抽取成独立的函数。结构清晰，意图明白，对于编写和阅读都是极大的便利。不过，在某些特殊的情况下，还可以使用其它更为优雅的写法，比如我们之前分享过的《不使用 if-elif 语句，如何优雅地判断某个数字所属的等级？》，以及今天要分享的这篇文章，可以开阔一下代码编写思路。

今天在 Github 阅读EdgeDB[1]的代码，发现它在处理大量if...elif...else判断的时候，使用了一个非常巧妙的装饰器。我们来看看这个方法具体是什么样的。

正好今天是双十一，假设我们要做一个功能，根据用户的等级判断他可以获得的折扣。常规的if ... elif...写法是这样的：

def get_discount(level):
    if level == 1:
        "大量计算代码"
        discount = 0.1
    elif level == 2:
        "大量计算代码"
        discount = 0.2
    elif level == 3:
        discount = 0.3
    elif level == 4:
        discount = 0.4
    elif level == 5:
        discount = 0.5
    elif level == 6:
        discount = 3 + 2 - 5 * 0.1
    else:
         return '等级错误'
    return discount

大家都知道，这样大量的if ... elif...代码非常难看，也很难维护。并且每个 if 的内部有很多代码。这个函数就会被拉得非常长。

有一些同学知道，可以使用字典来改写这个太长的 if 判断：

def parse_level_1():
    "大量计算代码"
    discount = 0.1
    return discount

def parse_level_2():
    "大量计算代码"
    discount = 0.2
    return discount

def parse_level_3():
    "大量计算代码"
    discount = 0.3
    return discount

def parse_level_4():
    "大量计算代码"
    discount = 0.4
    return discount

def parse_level_5():
    "大量计算代码"
    discount = 0.5
    return discount

def parse_level_6():
    "大量计算代码"
    discount = 3 + 2 - 5 * 0.1
    return discount

discount_map = {
 1: parse_level_1,
  2: parse_level_2,
  3: parse_level_3,
  4: parse_level_4,
  5: parse_level_5,
  6: parse_level_6,
}

discount = discount_map.get(level, '等级错误')

但今天我学到的这个方法，比用字典更简单。我们先来看它的效果：

@value_dispatch
def get_discount(level):
    return '等级错误'

@get_discount.register(1)
def parse_level_1(level):
    "大量计算代码"
    discount = 0.1
    return discount

@get_discount.register(2)
def parse_level_2(level):
    "大量计算代码"
    discount = 0.2
    return discount

@get_discount.register(3)
def parse_level_3(level):
    "大量计算代码"
    discount = 0.3
    return discount

@get_discount.register(4)
def parse_level_4(level):
    "大量计算代码"
    discount = 0.4
    return discount

@get_discount.register(5)
def parse_level_5(level):
    "大量计算代码"
    discount = 0.5
    return discount

@get_discount.register(6)
def parse_level_1(level):
    "大量计算代码"
    discount = 3 + 2 - 5 * 0.1
    return discount


discount = get_discount(3)
print(f'等级3的用户，获得的折扣是：{discount}')

运行效果如下图所示：

这样写，比用字典的方式更直观，比直接用if ... elif...更简洁。

那么，这个装饰器value_dispatch是怎么实现的呢？密码就藏在这个开源项目EdgeDB的源代码[2]中，核心代码只有20多行：

并且，还能够实现或查询。例如用户等级为2或者3的时候，折扣都是0.2，那么代码可以写成：

@get_discount.register(2)
@get_discount.register(3)
def parse_level_2(level):
    "大量计算代码"
    discount = 0.2
    return discount

运行效果如下图所示：

它这个代码目前只能实现相等的查询。但其实只要对这个代码稍作修改，我们就能实现大于、小于、大于等于、小于等于、不等于、in等等判断。如果大家有兴趣的话，请在文章下面留言，我们明天就来说说怎么对这个代码进行改造，实现更多的逻辑判断。

参考文献

[1] EdgeDB: https://github.com/edgedb/edgedb

[2] 源代码: https://github.com/edgedb/edgedb/blob/master/edb/common/value_dispatch.py

Python猫技术交流群开放啦！群里既有国内一二线大厂在职员工，也有国内外高校在读学生，既有十多年码龄的编程老鸟，也有中小学刚刚入门的新人，学习氛围良好！想入群的同学，请在公号内回复『交流群』，获取猫哥的微信（谢绝广告党，非诚勿扰！）~

还不过瘾？试试它们

▲终于懂了：协程思想的起源与发展

▲Python爬虫最强组合httpx+parsel测评

▲收藏了 30+ 实用的 Python 办公自动化库！

▲通过“四不要”，掌握 Python 的 Lambda 函数

▲OrderedDict 是如何保证 Key 的插入顺序的？

▲Python猫荐书系列之五：Python高性能编程

如果你觉得本文有帮助

请慷慨分享和点赞，感谢啦！