再見 if…elif… ，使用 Python 裝飾器輕鬆拿下

大家好，今天在 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/edged

最後祝大家天天進步！！學習Python最重要的就是心態。我們在學習過程中必然會遇到很多難題，可能自己想破腦袋都無法解決。這都是正常的，千萬別急著否定自己，懷疑自己。如果大家在剛開始學習中遇到困難，想找一個python學習交流環境，可以加入我們，領取學習資料，一起討論。