一日一技：使用装饰器简化大量 if…elif…代码

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

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

1. def get_discount(level):
2. if level == 1:
3. "大量计算代码"
4. discount = 0.1
5. elif level == 2:
6. "大量计算代码"
7. discount = 0.2
8. elif level == 3:
9. discount = 0.3
10. elif level == 4:
11. discount = 0.4
12. elif level == 5:
13. discount = 0.5
14. elif level == 6:
15. discount = 3 + 2 - 5 * 0.1
16. else:
17. return '等级错误'
18. return discount

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

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

1. def parse_level_1():
2. "大量计算代码"
3. discount = 0.1
4. return discount
6. def parse_level_2():
7. "大量计算代码"
8. discount = 0.2
9. return discount
11. def parse_level_3():
12. "大量计算代码"
13. discount = 0.3
14. return discount
16. def parse_level_4():
17. "大量计算代码"
18. discount = 0.4
19. return discount
21. def parse_level_5():
22. "大量计算代码"
23. discount = 0.5
24. return discount
26. def parse_level_6():
27. "大量计算代码"
28. discount = 3 + 2 - 5 * 0.1
29. return discount
31. discount_map = {
32. 1: parse_level_1,
33. 2: parse_level_2,
34. 3: parse_level_3,
35. 4: parse_level_4,
36. 5: parse_level_5,
37. 6: parse_level_6,
38. }
40. discount = discount_map.get(level, '等级错误')

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

1. @value_dispatch
2. def get_discount(level):
3. return '等级错误'
5. @get_discount.register(1)
6. def parse_level_1(level):
7. "大量计算代码"
8. discount = 0.1
9. return discount
11. @get_discount.register(2)
12. def parse_level_2(level):
13. "大量计算代码"
14. discount = 0.2
15. return discount
17. @get_discount.register(3)
18. def parse_level_3(level):
19. "大量计算代码"
20. discount = 0.3
21. return discount
23. @get_discount.register(4)
24. def parse_level_4(level):
25. "大量计算代码"
26. discount = 0.4
27. return discount
29. @get_discount.register(5)
30. def parse_level_5(level):
31. "大量计算代码"
32. discount = 0.5
33. return discount
35. @get_discount.register(6)
36. def parse_level_1(level):
37. "大量计算代码"
38. discount = 3 + 2 - 5 * 0.1
39. return discount
42. discount = get_discount(3)
43. print(f'等级3的用户，获得的折扣是：{discount}')

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

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

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

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

运行效果如下图所示：

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

参考文献

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

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

原文链接：https://mp.weixin.qq.com/s/eGji_6PBkf-vK2SQPB0WTQ