[筆記] Python - 應用 decorator 實作函式開關

　　最近為了課程撰寫遊戲的 server 端 (專案於 Github)，主要提供遊戲資訊跟進程控制的功能，client 端是一群擁有通訊功能的迷宮車，兩邊靠既定的指令溝通。在設計遊戲功能時，有些指令希望只在遊戲開始後才有作用，但又不可能讓 client 端乖乖在遊戲開始後才發送指令。最直觀的作法就是在每個處理指令的函式中檢查遊戲是否開始：

def game_command(...):
    if not game_is_started:
        return
    # Other jobs

但很快發現充滿重複的程式碼，而且影響程式碼的簡潔與維護的不方便。於是上網查找有無類似 C# 的 function attribute 的寫法，只要在函式的定義前加上像是 game_started 的 attribute，就可以控制該函式的行為。於是找到了 decorator (裝飾器)。

　　本篇主要介紹我如何在專案中使用 decorator，先簡單介紹 decorator，接著是本篇主題 ─ 利用 decorator 來實作函式開關，與在類別中使用 decorator，最後是如何進一步強化 decorator 的能力。以下講解皆使用 python 3。

Decorator

---

　　Decorator 是 python 的語法糖 (syntactic sugar)，會為函式加上一層 wrapper (包裹器)，在執行函式之前能在 wrapper 中作額外的處理，像是計時、記錄函式使用情形。Decorator 也是一個函式，以函式作為輸入，並會回傳一個新的函式。一個簡單的 decorator 如下：

def print_called(func):              # decorator 以函式作為輸入
    def wrapper(*args, **kargs):     # wrapper 實際提供額外處理，以 *args, **kargs 能夠接受任何數量的參數
        print(func.__name__)
        return func(*args, **kargs)  # 最後呼叫原本的函式，並保持其回傳值
    return wrapper                   # 回傳一個新的函式

這個簡單的 decorator 會印出要被呼叫的函式名稱，將 decorator 加到一個函式上：

@print_called
def add(a, b):
    return a + b

>>> add(3, 4)
add
7

一個函式被套用 decorator 後，被作了轉換，能讓原本的函式加上 wrapper，實際的呼叫過程為：

add = print_called(add)  # decorator 的重要概念
# => wrapper  因為 print_called 會回傳 wrapper 函式

>>> add.__name__
'wrapper'

原本的 add 函式被替換成 wrapper 這個函式。加上參數就不難理解 decorator 與內部的 wrapper 運作方式；

add(3, 4) = print_called(add)(3, 4)
# => wrapper(3, 4)

　　在上述例子中，有一個問題：那就是 add 函式原本的屬性消失了，add.__name__ 不是回傳 'add'，而是 'wrapper'。要保留被包覆的函式原本的屬性，就必須使用 @wraps decorator：

from functools import wraps

def print_called(func):
    @wraps(func)                     # 加在 wrapper 之前
    def wrapper(*args, **kargs):
        print(func.__name__)
        return func(*args, **kargs)
    return wrapper 

>>> add.__name__
'add'

除了函式原本的名稱，原本的 __doc__ 與 __annotations__ 也會被保留下來。另外，@wraps 提供可以直接取用被裝飾的函式的重要功能：

>>> add.__wrapped__(3, 4)
7

在撰寫 decorator 的時候都應該要加上 @wraps decorator。

用 decorator 製作函式開關

---

　　在這個專案中，因為無法限制玩家甚麼時候發出指令，所以必須在指令對應的 callback function 中加上 game_started 的 decorator：

def game_started(func):
    @wraps(func)
    def wrapper(*args, **kargs):
        if is_game_started:
            return func(*args, **kargs)
        else:
            pass
    return wrapper

這個 decorator 會去檢查對應的 flag，即遊戲是否開始，如果已經開始，就執行原本的函式，如果還沒有開始，就不會作任何事情。套用到目標函式：

@game_started
def game_command():
    print("game_command")

>>> is_game_started = True
>>> game_command()
game_command
>>> is_game_started = False
>>> game_command()
>>>

可以發現在 is_game_started 為 True 時，會執行 game_command 中的程式，但為 False 時，就不會執行，如此一來就可以達到函式開關的功能。除此之外，對於其他類似的函式，不需要複製檢查 flag 的程式碼，只要在該函式加上一樣 decorator，就可以達成一樣的效果。

@game_started
def game_another_command(...):
    ...

@game_started
def game_other_commands(...):
    ...

而且使用 @game_started 很容易讓人理解，這個函式在遊戲開始時才有作用，增加程式的可讀性與維護性。
再利用 is_game_started 製作一個 @game_stopped decorator，來讓函式只有在遊戲停止時才有作用：

def game_stopped(func):
    @wraps(func)
    def wrapper(*args, **kargs):
        if not is_game_started:
            return func(*args, **kargs)
        else:
            pass
    return wrapper

@game_stopped
def game_start(...):
    ...

@game_started
def game_stop(...):
    ...

在類別中使用 decorator

---

　　在專案中，遊戲相關的函式都定義在 GameCore 類別裡，包含遊戲是否開始的 flag。與前面不同的是，decorator 需要能取得類別物件的屬性，以 @game_started 為例：

class GameCore:
    def __init__(self):
        self.is_game_started = False

    def game_started(func):                       # 注意這裏不需加 self
        @wraps(func)
        def wrapper(self, *args, **kargs):        # 在 wrapper 中第一個參數必須為接受類別物件的參數
            if self.is_game_started:
                return func(self, *args, **kargs) # 使用原本函式也要傳入類別物件
            else:
                pass
        return wrapper

    @game_started
    def game_command(self):
        print("game_command")

>>> game_core = GameCore()
>>> game_core.game_command()
>>> game_core.is_game_started = True
>>> game_core.game_command()
game_command

self 加在 wrapper function 而不是在 decorator 上，一樣可以由 decorator 的運作推得被轉換後的函式：

GameCore.game_command = GameCore.game_started(GameCore.game_command)
# => wrapper

而類別物件透過 . (dot) 來呼叫類別函式是語法糖，實際運作方式為：

game_core.game_command() = GameCore.game_command(game_core) # 類別函式呼叫的語法糖
= GameCore.game_started(GameCore.game_command)(game_core)
# => wrapper(game_core)

可以發現，即使 decorator 需要用到類別物件的屬性，也不一定要宣告在類別之中，傳進的類別物件可以是任何類別，只要擁有 is_game_started 屬性。但如果需要使用到某個類別的屬性時，decorator 還是宣告在該類別之中比較好。
　　如果要在繼承類別中使用父類別的 decorator，要注意以下幾點：
1. 使用父類別的 decorator，必須使用父類別的名稱，無法使用 super()

class CustomGameCore(GameCore):
    def __init__(self):
        super().__init__()

    @GameCore.game_started
    def custom_game_command(self):
        print("custom_game_command")

>>> game_core = CustomGameCore()
>>> game_core.custom_game_command()
>>> game_core.is_game_started = True
>>> game_core.custom_game_command()
custom_game_command

2. 覆寫 (override) 父類別套有 decorator 的函式，不會帶有該 decorator 的特性。與
3. 使用父類別套有 decorator 函式，會保留該 decoator 的特性

class CustomGameCore(GameCore):
    def __init__(self):
        super().__init__()

    def game_command(self):
        super().game_command()
        print("custom_game_command")

>>> game_core = CustomGameCore()
>>> game_core.game_command()         # 第二點
custom_game_command
>>> game_core.is_game_started = True # 第三點
>>> game_core.game_command()
game_command
custom_game_command

因此，如果要覆寫父類別的函式，也要一併加上該函式套用的 decorator，才能保持一致的特性：

class CustomGameCore(GameCore):
    def __init__(self):
        super().__init__()

    @GameCore.game_started
    def game_command(self):
        super().game_command()
        print("custom_game_command")

>>> game_core = CustomGameCore()
>>> game_core.game_command()
>>> game_core.is_game_started = True
>>> game_core.game_command()
game_command
custom_game_command

強化 decorator 的能力：帶有引數的 decorator

---

　　在專案的測試階段中，有些功能希望可以不必在遊戲開始時才能使用，也就是希望 decorator 能夠有彈性，在測試狀態下可以解除限制，可以以 @game_started(disabled_in_test = True) 裝飾在需要解除限制的函式上：

class GameCore:
    def __init__(self):
        self.test_mode = False        # test_mode 也有可能是一個 global variable
        self.is_game_started = False

    def game_started(disabled_in_test = False):
        def decorator(func):

           @wraps(func)
           def wrapper(self, *args, **kargs):
              if (disabled_in_test and self.test_mode) or \
                 self.is_game_started:
                 return func(self, *args, **kargs)
              else:
                 pass
           return wrapper

        return decorator

    @game_started()            # 要注意即使不需要額外設定，也要加上 ()
    def game_command(self):
        print("game_command")

    @game_started(disabled_in_test = True)
    def another_game_command(self):
        print("another_game_command")

>>> game_core = GameCore()
>>> game_core.test_mode = True
>>> game_core.game_command()
>>> game_core.another_game_command()
another_game_command

這個 decorator 利用使用者指定的 flag 並配合 test_mode 變數，來決定是否要執行被包覆的函式。寫出指定的參數是為了增加可讀性。
　　整個 decorator 多加一層的原因是為了讓 decorator 能接受引數：

another_game_command = game_started(disabled_in_test = True)(another_game_command)
# => decorator(another_game_command)
# => wrapper

要注意即使不需要額外指定，只要 decorator 帶有引數，就還是要加上 ()：

...

@game_started     # 沒有加上 ()
def game_command(self):
    print("game_command")

...

>>> game_core.foo()
<function GameCore.game_started.<locals>.decorator.<locals>.wrapper at 0x0000000002182E18>

沒有加上 () 的話，decorator 還是以 func = decorator(func) 去處理。雖然不會報錯，但是每層函式傳入的參數都不是預期的參數，執行結果也不是預期的結果。

game_command = game_started(disabled_in_test = game_command)
# => decorator

game_core.game_command() = GameCore.game_command(game_core)
# => decorator(func = game_core) 只會回傳 wrapper 的位址

如果有加上 () 的話，：

game_command = game_started()(game_command) # disabled_in_test 是選擇性引數
# => decorator(game_command)
# => wrapper

game_core.game_command() = GameCore.game_command(game_core)
# => wrapper(game_core)

附錄：製作帶有選擇性引數的 decorator (不需額外加 ())

---

　　會把這一段寫在附錄的原因是這種作法不能套用在類別上，會有 NameError 的問題。以下的例子皆不是在類別之中，is_game_started 與 test_mode 都是 global variable，製作一個帶有選擇性引數的 decorator：

from functools import wraps, partial

is_game_started = False
test_mode = False

def game_started(func = None, *, disabled_in_test = False):
    if func is None:
        return partial(game_started, disabled_in_test = disabled_in_test)

    @wraps(func)
    def wrapper(*args, **kargs):
        if (disabled_in_test and test_mode) or \
            is_game_started:
            return func(*args, **kargs)
        else:
            pass

    return wrapper

@game_started
def game_command():
    print("game_command")

@game_started(disable_in_test = True)
def another_game_command():
    print("another_game_command")

>>> game_command()
>>> another_game_command()
>>> test_mode = True
>>> game_command()
>>> another_game_command()
another_game_command

能夠不加 () 主要靠 functools 的 partial，他能夠固定某些引數的值，回傳一個所需引數較少的函式：

def add(a, b, c):
    return a + b + c

foo = partial(add, c = 3)
# => foo(a, b) = add(a, b, c = 3)

foo_2 = partial(add, 2, c = 3)
# => foo_2(b) = add(2, b, c = 3)

>>> foo(1, 2)
6
>>> foo_2(4)
9

所以當 decorator 有指定引數時：

another_game_command = game_started(disabled_in_test = True)(another_game_command)
# 第一次呼叫 game_started，func 為 None，所以 partial 回傳已經指定好 disabled_in_test 的 game_started 函式
= game_started(another_game_command, disabled_in_test = True)
# => wrapper

如果沒有指定引數的話：

game_command = game_started(game_command) # func = game_command
# => wrapper

game_started 函式的引數定義中有個 * (匿名的位置引數)，是為了強迫使用者指定在這之後的引數名稱的技巧。如果使用者沒有指定名稱的話，除了前面的位子參數 (positional arguments) 之外，其餘參數都會被指定到 * 中。在這裡是為了避免使用者出現 @game_started(True) 這樣的用法，因為如此一來 func 就會被指定為 True 了。
　　至於不能使用在類別定義中的原因為，還在讀取類別定義時，有套用 decorator 的函式會先被執行一次轉換，而此時如果取用類別名稱 (也就是 partial 中的目標函式為 GameCore.game_started) 會出現 NameError: name 'GameCore' is not defined，因為該類別還在讀取定義中。

參考資料

---

• Python 錦囊妙計 (Python Cookbook), O'Reilly
• Forced naming of parameters in python, StackOverflow
• Guild to python function decorators, The code ship