日常 Python：如何使用 assert 提高代碼品質，讓 bug 無所遁形

#python
五倍技術部
技術文章
日常 Python：如何使用 assert 提高代碼品質，讓 bug 無所遁形

日常 Python：如何使用 assert 提高代碼品質，讓 Bug 無所遁形

在開始講如何使用 Assertion 之前，我們先來看看一段程式碼：

def apply_discount(product: dict, discount: float) -> int:
    price = int(product["price"] * (1.0 - discount))
    return price

這段程式碼的用途是計算商品折扣後的價格，它可以接收字典 product 跟浮點數 discount 兩個參數，所以假設我有一個商品是毛巾，售價為 2000 元，我想要打 75 折，我可以這樣寫：

product = {
    "name": "towel",
    "price": 2000,
}

discount = 0.25

apply_discount(product, discount) # 這條毛巾打完折的價格是 1500

那假設今天店員在上架 1000 元的帽子時，本來打算打八折（也就是 discount = 0.2），但一個不小心手滑，將 0.2 按成了 2，此時再看看程式碼：

product = {
    "name": "hat",
    "price": 1000,
}

discount = 2

apply_discount(product, discount) # 這條毛巾打完折的價格是 -1000

(買一頂賺一頂)

顯然，價格不可能為負數，否則就會導致商家虧損，就算打折打到老闆要跳樓，return 回傳折扣後的價格最少就是 0 元，於是身為 Python 初學者的我們，使用 if / else 將 apply_discount 做了一點優化：

def apply_discount(product: dict, discount: float) -> int:
    price = int(product["price"] * (1.0 - discount))

    if price >= 0:
        return price

    else:
        raise ValueError("折扣後價格不能小於 0")

product = {
    "name": "hat",
    "price": 1000,
}

discount = 2

apply_discount(product, discount)

乍看之下好像沒什問題，但讓我們重新思考一下 if / else 所要表達的語意是什麼？

通常我們會用 if / else 表達各種狀況到流程控制，稍微翻譯一下大概是：「if 是 A 狀況，則是 A 結果，else 其他的狀況，通通都是 B 結果」，換句話說 if 跟 else 都是有可能發生的狀況，所以我們才需要思考對應的結果。

回到 apply_discount 商品折扣的範例，對我們而言，最後打折完的售價無論如何都要大於等於 0，也就是說如果發生小於 0 的狀況，程式應該立即中斷，並回報錯誤：「這邊錯了，不可能小於 0！」，此時的「小於 0」應該視為一個錯誤（bug）而不是一種狀況（condition），既然如此 if / else 可能就不是最好的選擇。

使用 assert 來解決這個問題

Assertion（斷言）聽翻譯就有種「鐵口直斷」的感覺，可以把它想像成「只有我說的這種可能，其他都是錯的」，如果是 apply_discount 這個例子，那就是「結果只可能大於等於 0，不然一定是錯的」，換成程式碼的話，可以這樣改寫 apply_discount：

def apply_discount(product: dict, discount: float) -> int:
    assert 0.0 <= discount <= 1.0
    price = int(product["price"] * (1.0 - discount))
    return price

所以當我們再次執行下面的程式碼時：

product = {
    "name": "hat",
    "price": 1000,
}

discount = 2

apply_discount(product, discount)

會發生這樣的 AssertionError：

Traceback (most recent call last):
  File "/Users/chienchuanw/Documents/django-template/tempCodeRunnerFile.python", line 14, in <module>
    apply_discount(product, discount)
    ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^
  File "/Users/chienchuanw/Documents/django-template/tempCodeRunnerFile.python", line 2, in apply_discount
    assert 0.0 <= discount <= 1.0
           ^^^^^^^^^^^^^^^^^^^^^^
AssertionError

使用 Assertion 有什麼好處？

相較於先前的 if / else 不僅會完整地運行完程式碼，就算我們看到了 "出問題了，結果不應該小於 0"，我們也只能知道有地方不對，但無法迅速找出錯誤發生的地方，但 Assertion 就不一樣了，終端機不僅噴錯，還很明確地說明：

  • 在第十四行（line 14），使用 apply_discount(product, discount) 時出錯。
  • 這個錯誤可以追溯到第二行（line 2） apply_discount 內部裡，不符合 0 <= price <= product["price"] 這個斷言。

所以我們可以很快地得知，是因為 price 計算出來的結果不符合規定，因此引發了錯誤。

這邊如果想讓使用者能更快的理解 0 <= price <= product["price"] 所代表的意涵，還可以更一步地善用 assert 的警示文字：

def apply_discount(product: dict, discount: float) -> int:
    assert 0.0 <= discount <= 1.0, "折扣必須介於 0.0 和 1.0 之間"
    price = int(product["price"] * (1.0 - discount))
    return price


product = {
    "name": "hat",
    "price": 1000,
}

discount = 2

apply_discount(product, discount)

這樣錯誤訊息也會顯示原因，而不是單純拋出 AssertionError

Traceback (most recent call last):
  File "/Users/chienchuanw/Documents/django-template/tempCodeRunnerFile.python", line 14, in <module>
    apply_discount(product, discount)
    ~~~~~~~~~~~~~~^^^^^^^^^^^^^^^^^^^
  File "/Users/chienchuanw/Documents/django-template/tempCodeRunnerFile.python", line 2, in apply_discount
    assert 0.0 <= discount <= 1.0, "折扣必須介於 0.0 和 1.0 之間"
           ^^^^^^^^^^^^^^^^^^^^^^
AssertionError: 折扣必須介於 0.0 和 1.0 之間

所以我們可以將 assert 當作是開發階段的錯誤檢查工具，它的主要用途在於不變性檢查（invariant checking），也就是確保某些條件在程式執行期間始終成立。

Assertion 也有踩雷的時候

好的你現在學會了 Assertion，所以迫不及待地在程式碼上到處 assert，但不要以為 assert 百利而無一害，錯誤地使用 assert 不僅沒能除錯，甚至還會程式產生安全漏洞，進一步地被惡意攻擊。

雷點一：不應該使用 assert 進行資料驗證

def delete_product(prod_id: int, user: dict) -> None:
    assert user.is_admin(), "必須是 admin"
    assert store.has_product(prod_id), "未知的 product"
    store.get_product(prod_id).delete()

根據這段 delete_product （刪除商品）想表達的意思，應該是：

  • 斷言 user.is_admin() 必須為 True，也就是說 user 必須是 admin。
  • 斷言 store.has_product(prod_id) 必須為 True，也就是說 store 裡面必須有這個 prod_id。
  • 通過前面兩個斷言後，store 會刪除這 prod_id。

乍看之下好像沒什麼問題，但如果我們在某些部署環境（如特定的生產環境），使用優化模式（optimize mode）執行 Python 程式碼時，此時所有的 assert 語句都會被忽略，也就是說當我們執行 delete_product 時，不管「user 是不是 admin」、「store 有沒有 prodid」，都會執行 `store.getproduct(prod_id).delete()`，這樣會發生兩個問題：

我們可以使用 $ python -O 進入 Optimize mode，或是使用 $ python -OO 進入 Extra Optimize mode，兩種優化模式都會忽略 assert

  1. 不是 admin 的 user 也可以任意刪除 store 裡的任意 prod_id 商品。
  2. 有心人可以透過大量發送無效（不存在）的 prod_id 影響程式的效能，甚至可能導致伺服器崩潰，也就是所謂的 DDoS (Distributed Denial of Service) 攻擊。

所以我們可以這樣改寫程式碼，避免因為優化模式所帶來的副作用，針對每一個步驟的異常錯誤，拋出對應處理與訊息：

def delete_product(prod_id: int, user) -> None:
    if not user.is_admin():
        raise PermissionError("Must be admin")

    if not store.has_product(prod_id):
        raise ValueError("Unknown product")

    product = store.get_product(prod_id)
    if product is None:
        raise ValueError("獲取 product 失敗，可能已被刪除")

    try:
        product.delete()
    except Exception as e:
        raise RuntimeError(f"刪除 product 失敗: {e}")

既然優化模式會忽略 assert，那這時候我們再回去看稍早講的 apply_discount 程式碼，不禁會納悶優化模式下，assert 0 <= price <= product["price"] 不就也沒作用了嗎？price 最後仍舊有可能有可能會出現小於 0 的狀況。

def apply_discount(product: dict, discount: float) -> int:
    assert 0.0 <= discount <= 1.0, "折扣必須介於 0.0 和 1.0 之間" # 這段在優化模式下，會被忽略
    price = int(product["price"] * (1.0 - discount))
    return price

假設 apply_discount 是重要的商業計算邏輯，那確實不應該使用 assert，避免遇到負價格出現的可能，進而導致整個商店系統的崩壞，所以比較好的做法也是透過 if 檢查後，使用 raise 拋出錯誤異常：

def apply_discount(product: dict, discount: float) -> int:
    assert 0.0 <= discount <= 1.0, "折扣必須介於 0.0 和 1.0 之間"

    price = int(product["price"] * (1.0 - discount))

    if price < 0:
        raise ValueError("折扣後價格不可為負數")

    return price

(開發的時候都沒問題，正式上線也不會有問題(吧？))

雷點二：不小心寫出永遠不會出錯的 assert

先看看這段程式碼：

a = 1
b = 2

assert (a >= b, "a 應該要大於 b")
print("通過測試")

單就這段程式碼，我們可以猜測其中想要表達的意思應該是：「使用 assert 斷言 a 要大於 b，如果不是的話，則會拋出 AssertionError」，所以上面程式碼應該要跳出錯誤才對，但怎麼結果不僅沒跳出錯誤，還順順印出了"通過測試"的字串了呢？

因為在 Python 中，assert 並不是一個函數（function），而是一個關鍵字（keyword）！

正確的 assert 語法是這樣：

assert 判斷的結果（boolean）, 顯示的錯誤訊息

讓我們逐步拆解 assert (a >= b, "a 應該要大於 b") 所代表的意思：

  • assert (a >= b, "a 應該要大於 b") 這一段的括號 () 實際上是把 (a >= b, "a 應該要大於 b") 變成了一個元組（tuple），元組裡的第一個值（a >= b）會得到一個布林值（boolean），第二個值（"a 應該要大於 b"）會得到字串（string）。
  • a >= b 經過程式計算後會得到 False，所以程式碼會變成 assert (False, "a 應該要大於 b")。
  • 由於 assert 會針對判斷的結果（boolean）為 False，才會顯示錯誤訊息，所以此處 Python 會將元組 (False, "a 應該要大於 b") 轉換為 True。
  • 程式碼最後會變成 assert True，而 assert True 永遠不會觸發到 AssertionError。

因此如果要讓程式碼發生預期的檢查，應該要改為：

a = 1
b = 2

assert a >= b, "a 應該要大於 b"
print("通過測試")

(少一個括號 () 差很多啊！)

小結

Assertion 作為一個 debug 工具可以幫助我們快速找出程式碼問題所在之處，然而也不是任何狀況使用 assert 都是最好的選擇，通常會在這些情境下使用 assert：

  • 開發、測試環境
  • 為了確保內部邏輯正確
  • 單元測試
  • 不影響正式環境

因此 assert 主要用於內部檢查與測試，協助工程師盡快找出問題之所在，而不能用來處理執行時所發生的錯誤，或是權限或商業邏輯檢查，如果會因為 assert 被忽略後而影響程式本來的行為，那麼就應該改用 if / raise Exception，而攸關安全性、資料驗證、商業邏輯等關鍵代碼，更不應該使用 assert，以防造成不可預想的後果。

如果本文章對你有幫助，歡迎按讚或者留言討論 ヽ(●´∀`●)ﾉ
本文同步發佈於作者的 Medium 網站。