DEV Community

qing
qing

Posted on

3 Ways Enums Boost Python Code

Python enum: Better Constants and State Machines

enum replaces magic strings and numbers with readable, type-safe constants.

Basic Enum

from enum import Enum, auto

class Color(Enum):
    RED = 1
    GREEN = 2
    BLUE = 3

# Usage
color = Color.RED
print(color)         # Color.RED
print(color.name)    # RED
print(color.value)   # 1
print(color == Color.RED)  # True

# Iteration
for c in Color:
    print(c.name, c.value)
Enter fullscreen mode Exit fullscreen mode

auto() — Automatic Values

from enum import Enum, auto

class Direction(Enum):
    NORTH = auto()  # 1
    SOUTH = auto()  # 2
    EAST = auto()   # 3
    WEST = auto()   # 4

class Status(Enum):
    PENDING = auto()
    RUNNING = auto()
    SUCCESS = auto()
    FAILED = auto()

def process(status: Status) -> str:
    match status:
        case Status.PENDING:  return "Waiting..."
        case Status.RUNNING:  return "In progress"
        case Status.SUCCESS:  return "Done!"
        case Status.FAILED:   return "Error!"
Enter fullscreen mode Exit fullscreen mode

IntEnum for Comparisons

from enum import IntEnum

class Priority(IntEnum):
    LOW = 1
    MEDIUM = 2
    HIGH = 3
    CRITICAL = 4

# Can compare with integers
task_priority = Priority.HIGH
if task_priority >= Priority.HIGH:
    send_alert()

# Sort by priority
tasks = [
    ("fix bug", Priority.HIGH),
    ("write docs", Priority.LOW),
    ("deploy", Priority.CRITICAL),
]
sorted_tasks = sorted(tasks, key=lambda t: t[1], reverse=True)
Enter fullscreen mode Exit fullscreen mode

Flag for Bitwise Operations

from enum import Flag, auto

class Permission(Flag):
    READ = auto()     # 1
    WRITE = auto()    # 2
    EXECUTE = auto()  # 4

    # Combinations
    READ_WRITE = READ | WRITE
    ALL = READ | WRITE | EXECUTE

# Usage
user_perms = Permission.READ | Permission.WRITE
print(Permission.READ in user_perms)  # True
print(Permission.EXECUTE in user_perms)  # False

# Check permissions
def can_write(perms: Permission) -> bool:
    return Permission.WRITE in perms
Enter fullscreen mode Exit fullscreen mode

State Machine with Enum

from enum import Enum

class OrderState(Enum):
    CREATED = "created"
    PAID = "paid"
    SHIPPED = "shipped"
    DELIVERED = "delivered"
    CANCELLED = "cancelled"

VALID_TRANSITIONS = {
    OrderState.CREATED: [OrderState.PAID, OrderState.CANCELLED],
    OrderState.PAID: [OrderState.SHIPPED, OrderState.CANCELLED],
    OrderState.SHIPPED: [OrderState.DELIVERED],
    OrderState.DELIVERED: [],
    OrderState.CANCELLED: [],
}

class Order:
    def __init__(self):
        self.state = OrderState.CREATED

    def transition(self, new_state: OrderState):
        if new_state not in VALID_TRANSITIONS[self.state]:
            raise ValueError(f"Cannot go from {self.state} to {new_state}")
        self.state = new_state
        return self
Enter fullscreen mode Exit fullscreen mode

Follow me for more Python tips! 🐍


喜欢这篇文章?关注获取更多Python自动化内容!

Top comments (0)