Enum(列舉)用於定義一組具名的常數值。當你有一組固定的選項時,使用 Enum 比使用字串或數字更安全、更易讀。

為什麼使用 Enum?

使用字串的問題

1# 使用字串:容易打錯
2def handle_decision(decision: str) -> None:
3    if decision == "alow":  # 打錯了!應該是 "allow"
4        allow_action()
5    elif decision == "deny":
6        deny_action()
7
8# 沒有 IDE 自動完成
9# 沒有型別檢查

使用 Enum 的好處

 1from enum import Enum
 2
 3class Decision(Enum):
 4    ALLOW = "allow"
 5    DENY = "deny"
 6    ASK = "ask"
 7
 8def handle_decision(decision: Decision) -> None:
 9    if decision == Decision.ALLOW:  # IDE 會自動完成
10        allow_action()
11    elif decision == Decision.DENY:
12        deny_action()
13
14# 打錯會有 AttributeError
15# Decision.ALOW  # 錯誤!

基本用法

定義 Enum

 1from enum import Enum
 2
 3class Color(Enum):
 4    RED = 1
 5    GREEN = 2
 6    BLUE = 3
 7
 8# 存取
 9print(Color.RED)        # Color.RED
10print(Color.RED.name)   # 'RED'
11print(Color.RED.value)  # 1

字串值的 Enum

 1from enum import Enum
 2
 3class LogLevel(Enum):
 4    DEBUG = "debug"
 5    INFO = "info"
 6    WARNING = "warning"
 7    ERROR = "error"
 8
 9# 從值建立
10level = LogLevel("info")
11print(level)  # LogLevel.INFO

實際範例:Hook 決策

雖然 Hook 系統目前使用字串,但可以用 Enum 改善:

 1from enum import Enum
 2
 3class HookDecision(Enum):
 4    """Hook 決策類型"""
 5    ALLOW = "allow"
 6    DENY = "deny"
 7    ASK = "ask"
 8    BLOCK = "block"
 9
10class HookEventType(Enum):
11    """Hook 事件類型"""
12    PRE_TOOL_USE = "PreToolUse"
13    POST_TOOL_USE = "PostToolUse"
14    STOP = "Stop"
15    SESSION_START = "SessionStart"
16
17# 使用
18def create_output(decision: HookDecision) -> dict:
19    return {
20        "hookSpecificOutput": {
21            "permissionDecision": decision.value
22        }
23    }

進階功能

auto() 自動值

1from enum import Enum, auto
2
3class Priority(Enum):
4    LOW = auto()     # 1
5    MEDIUM = auto()  # 2
6    HIGH = auto()    # 3

IntEnum

當需要 Enum 值可以用作整數時:

 1from enum import IntEnum
 2
 3class HttpStatus(IntEnum):
 4    OK = 200
 5    NOT_FOUND = 404
 6    SERVER_ERROR = 500
 7
 8# 可以直接比較整數
 9if response.status == HttpStatus.OK:
10    ...
11
12# 可以用於數學運算
13print(HttpStatus.OK + 1)  # 201

StrEnum (Python 3.11+)

1from enum import StrEnum
2
3class Color(StrEnum):
4    RED = "red"
5    GREEN = "green"
6    BLUE = "blue"
7
8# 可以直接當字串使用
9print(f"Color is {Color.RED}")  # "Color is red"

Flag(位元旗標)

 1from enum import Flag, auto
 2
 3class Permission(Flag):
 4    READ = auto()
 5    WRITE = auto()
 6    EXECUTE = auto()
 7
 8# 組合權限
 9user_perms = Permission.READ | Permission.WRITE
10
11# 檢查權限
12if Permission.READ in user_perms:
13    print("Can read")

迭代和比較

迭代所有成員

 1from enum import Enum
 2
 3class Status(Enum):
 4    PENDING = "pending"
 5    RUNNING = "running"
 6    COMPLETED = "completed"
 7
 8# 迭代
 9for status in Status:
10    print(f"{status.name}: {status.value}")
11
12# 取得所有值
13all_values = [s.value for s in Status]
14# ['pending', 'running', 'completed']

比較

 1from enum import Enum
 2
 3class Color(Enum):
 4    RED = 1
 5    GREEN = 2
 6
 7# Enum 成員是單例
 8Color.RED is Color.RED  # True
 9Color.RED == Color.RED  # True
10
11# 不能與原始值直接比較
12Color.RED == 1  # False
13Color.RED.value == 1  # True

從值建立 Enum

 1from enum import Enum
 2
 3class Status(Enum):
 4    ACTIVE = "active"
 5    INACTIVE = "inactive"
 6
 7# 從值建立
 8status = Status("active")  # Status.ACTIVE
 9
10# 從名稱建立
11status = Status["ACTIVE"]  # Status.ACTIVE
12
13# 安全地從值建立
14def get_status(value: str) -> Status:
15    try:
16        return Status(value)
17    except ValueError:
18        return Status.INACTIVE

實際應用模式

配置選項

 1from enum import Enum
 2
 3class OutputFormat(Enum):
 4    JSON = "json"
 5    YAML = "yaml"
 6    TEXT = "text"
 7
 8def export_data(data: dict, format: OutputFormat) -> str:
 9    if format == OutputFormat.JSON:
10        return json.dumps(data)
11    elif format == OutputFormat.YAML:
12        return yaml.dump(data)
13    else:
14        return str(data)

狀態機

 1from enum import Enum, auto
 2
 3class TaskState(Enum):
 4    PENDING = auto()
 5    RUNNING = auto()
 6    COMPLETED = auto()
 7    FAILED = auto()
 8
 9class Task:
10    def __init__(self):
11        self.state = TaskState.PENDING
12
13    def start(self):
14        if self.state != TaskState.PENDING:
15            raise ValueError("Cannot start non-pending task")
16        self.state = TaskState.RUNNING
17
18    def complete(self):
19        if self.state != TaskState.RUNNING:
20            raise ValueError("Cannot complete non-running task")
21        self.state = TaskState.COMPLETED

驗證等級

 1from enum import Enum
 2
 3class ValidationLevel(Enum):
 4    ERROR = "error"
 5    WARNING = "warning"
 6    INFO = "info"
 7
 8    def is_blocking(self) -> bool:
 9        """是否會阻止操作"""
10        return self == ValidationLevel.ERROR
11
12# 使用
13issue_level = ValidationLevel.WARNING
14if issue_level.is_blocking():
15    raise ValidationError()

最佳實踐

1. 使用全大寫命名成員

1class Status(Enum):
2    ACTIVE = "active"    # 好
3    active = "active"    # 不推薦

2. 為 Enum 添加方法

 1class Priority(Enum):
 2    LOW = 1
 3    MEDIUM = 2
 4    HIGH = 3
 5
 6    def is_urgent(self) -> bool:
 7        return self == Priority.HIGH
 8
 9    @classmethod
10    def from_string(cls, value: str) -> "Priority":
11        return cls[value.upper()]

3. 搭配型別提示使用

 1from enum import Enum
 2
 3class Color(Enum):
 4    RED = "red"
 5    BLUE = "blue"
 6
 7def paint(color: Color) -> None:  # 型別提示
 8    print(f"Painting with {color.value}")
 9
10paint(Color.RED)      # OK
11paint("red")          # 型別檢查會警告

思考題

  1. EnumIntEnum 有什麼區別?什麼時候用哪個?
  2. 為什麼 Color.RED == 1 會是 False
  3. 如何為 Enum 添加自訂方法?

實作練習

  1. 建立一個 HookType Enum,包含所有 Hook 事件類型
  2. 實作一個包含 is_valid() 方法的 FileExtension Enum
  3. 使用 Flag 實作一個權限系統

上一章:Dataclass 資料結構 下一模組:標準庫實戰