Python: defaultdict 的陷阱

最近壓力比較大，讓我出賣一下我的同事。

我的同事 L 最近要用 Python 寫一個函式 lookup()。它會拿三個參數：

d：一個 str 到 set 的 dict
x：當作 d 的鍵值的字串
y：字串或是 None

它的功能是：

如果 d 沒有 x，則 lookup() 應回傳 False。
如果 d 有 x 且 y 是 None，則 lookup() 應回傳 True。
如果 d 有 x 且 y 不是 None，則 lookup() 應回傳 y 是不是在 d[x] 之內。

這看起來很簡單，所以 L 就寫出以下的程式碼：

def lookup(d, x, y):
    try:
        s = d[x]
    except KeyError:
        return False
    if y is None:
        return True
    return y in s

看起來很合理。但是過幾天後，同事 L 發現這個程式的結果有問題。不過這份 code 看起來很完美，不可能有問題呀。

一個小時過後，同事 L 突然頓悟了：問題不在於 lookup() 而是在於 d 的型別。因為在建構 d 的時候，同事 L 是使用 collections.defaultdict(set) 作為型別：

def build_table(lines):
    d = collections.defaultdict(set)
    for line in lines:
        x, y = line.split(',', 1)
        d[x].add(y)
    return d

原本 lookup() 函式中的 d[x] 會馬上在 d 裡面加上一個空的 set 物件，然後回傳該 set 物件，所以 KeyError 例外永遠不會被拋出。

如果把上面 lookup() 的 d[x] 改為 d.get(x) 就沒問題了：

def lookup(d, x, y):
    s = d.get(x)
    if s is None:
        return False
    if y is None:
        return True
    return y in s

另一個方法是在 build_table() 回傳之前，把 defaultdict 複製成 dict：

def build_table(lines):
    d = collections.defaultdict(set)
    for line in lines:
        x, y = line.split(',', 1)
        d[x].add(y)
    return dict(d)

或者，乾脆不要使用 defaultdict：

def build_table(lines):
    d = {}
    for line in lines:
        x, y = line.split(',', 1)
        s = d.get(x)
        if s is not None:
            s.add(y)
        else:
            d[x] = set([y])
    return d

同事 L 暫時用第二個改法快速地 workaround。不過同事 L 仔細想想覺得或許改用 d.get(x) 會比較好。

但這衍生一個問題：要怎麼和 Python 的 EAFP 原則取得平衡？對於類似的查詢，我們是否都應該要保守地假設 d 有可能是 defaultdict，所以都只能使用 d.get(x) 查詢 dict-like 物件呢？

Python: defaultdict 的陷阱

羅根學習筆記