Pythonのデコレータの使い方入門: 関数を変更せず機能を拡張する

# デコレータの最小例
from typing import Callable

def hello_decorator(func: Callable) -> Callable:
    def wrapper():
        print("前処理: Hello")
        result = func()
        print("後処理: Bye")
        return result
    return wrapper

# @記法の利用
@hello_decorator
def greet():
    print("中身: greetが呼ばれました")

# 上の@は次と等価です:
# greet = hello_decorator(greet)

greet()

def notify(func):
    """呼び出しの前後でメッセージを表示する簡単なデコレータ"""
    def wrapper():
        print("[notify] start")
        value = func()
        print("[notify] end")
        return value
    return wrapper

@notify
def say():
    print("Hi!")

say()

from typing import Any, Callable

def debug_args(func: Callable[..., Any]) -> Callable[..., Any]:
    """引数と戻り値を表示するデコレータ"""
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        print(f"[debug_args] args={args}, kwargs={kwargs}")
        result = func(*args, **kwargs)
        print(f"[debug_args] result={result!r}")
        return result
    return wrapper

@debug_args
def add(a: int, b: int, scale: float = 1.0) -> float:
    return (a + b) * scale

add(2, 3, scale=0.5)

from functools import wraps
from typing import TypeVar, Callable, Any

F = TypeVar("F", bound=Callable[..., Any])

def preserve_meta(func: F) -> F:
    """wrapsを使ってメタデータを保持する例"""
    @wraps(func)  # これが重要です
    def wrapper(*args: Any, **kwargs: Any):
        """内側のdocstringは外に見えなくなる(元のdocstringが保持される)"""
        return func(*args, **kwargs)
    return wrapper  # type: ignore[return-value]

@preserve_meta
def area(radius: float) -> float:
    """半径から円の面積を計算します"""
    from math import pi
    return pi * radius * radius

print(area.__name__)   # 元の関数名が保たれる
print(area.__doc__)    # 元のdocstringが保たれる

from functools import wraps
from typing import Callable, Any

def repeat(times: int = 2) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """指定回数だけ関数を繰り返し実行するデコレータファクトリ"""
    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            result = None
            for i in range(times):
                print(f"[repeat] {i+1}/{times}")
                result = func(*args, **kwargs)
            return result
        return wrapper
    return decorator

@repeat(3)
def ping() -> str:
    print("pong")
    return "ok"

ping()

def A(func):
    def wrapper(*a, **k):
        print("A: before")
        r = func(*a, **k)
        print("A: after")
        return r
    return wrapper

def B(func):
    def wrapper(*a, **k):
        print("B: before")
        r = func(*a, **k)
        print("B: after")
        return r
    return wrapper

@A
@B
def work():
    print("do work")

work()

from functools import wraps
from typing import Any, Callable, Optional, Type, Tuple

def safe(default: Optional[Any] = None,
         suppress: Tuple[Type[BaseException], ...] = (Exception,)) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """
    例外をログし、指定の例外のみ抑制してdefaultを返す。
    それ以外は再送出します。
    """
    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            try:
                return func(*args, **kwargs)
            except suppress as e:
                print(f"[safe] suppressed: {e!r}")
                return default
        return wrapper
    return decorator

@safe(default=-1, suppress=(ZeroDivisionError,))
def div(a: int, b: int) -> float:
    return a / b

print(div(4, 2))  # 正常
print(div(1, 0))  # 例外を抑制しdefaultを返す

import time
from functools import wraps
from typing import Any, Callable

def timeit(label: str | None = None) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """実行時間を計測して表示するデコレータ"""
    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            name = label or func.__name__
            start = time.perf_counter()
            try:
                return func(*args, **kwargs)
            finally:
                dur = (time.perf_counter() - start) * 1000
                print(f"[timeit] {name}: {dur:.2f} ms")
        return wrapper
    return decorator

@timeit()
def slow(n: int) -> int:
    time.sleep(0.05)
    return n * 2

print(slow(5))

from dataclasses import dataclass
from functools import wraps
from typing import Callable, Any

@dataclass
class User:
    name: str
    is_admin: bool = False

def requires_admin(func: Callable[..., Any]) -> Callable[..., Any]:
    """第1引数userの権限をチェックするデコレータ"""
    @wraps(func)
    def wrapper(user: User, *args: Any, **kwargs: Any) -> Any:
        if not user.is_admin:
            raise PermissionError("admin権限が必要です")
        return func(user, *args, **kwargs)
    return wrapper

def validate_positive(func: Callable[..., Any]) -> Callable[..., Any]:
    """数値引数が正であることをバリデートする例"""
    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        for x in args:
            if isinstance(x, (int, float)) and x <= 0:
                raise ValueError("引数は正の数である必要があります")
        return func(*args, **kwargs)
    return wrapper

@requires_admin
@validate_positive
def create_quota(user: User, size_gb: int) -> str:
    return f"{user.name} に {size_gb}GB のクォータを作成しました"

admin = User("Alice", is_admin=True)
guest = User("Bob", is_admin=False)

print(create_quota(admin, 10))
# print(create_quota(guest, 10))  # PermissionError

from functools import lru_cache

@lru_cache(maxsize=128)
def fib(n: int) -> int:
    if n < 2:
        return n
    return fib(n-1) + fib(n-2)

print(fib(30))  # 高速に計算される(以降はキャッシュ)

import time
from functools import wraps
from typing import Callable, Any, Tuple, Type

def retry(retries: int = 3, delay: float = 0.1,
          exceptions: Tuple[Type[BaseException], ...] = (Exception,)) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    """一時的な失敗をリトライするデコレータ"""
    def decorator(func: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            last = None
            for attempt in range(1, retries + 1):
                try:
                    return func(*args, **kwargs)
                except exceptions as e:
                    last = e
                    print(f"[retry] attempt {attempt}/{retries} failed: {e!r}")
                    if attempt < retries:
                        time.sleep(delay)
            assert last is not None
            raise last
        return wrapper
    return decorator

count = {"n": 0}

@retry(retries=3, delay=0.05, exceptions=(RuntimeError,))
def flaky() -> str:
    count["n"] += 1
    if count["n"] < 3:
        raise RuntimeError("temporary failure")
    return "success"

print(flaky())

from functools import wraps
from typing import Any, Callable

def trace(func: Callable[..., Any]) -> Callable[..., Any]:
    """呼び出しをトレースするデコレータ"""
    @wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        print(f"[trace] {func.__qualname__} called")
        return func(*args, **kwargs)
    return wrapper

class Calculator:
    @trace
    def mul(self, a: int, b: int) -> int:
        return a * b

c = Calculator()
print(c.mul(3, 4))

from functools import wraps
from typing import Any, Callable

def tag(name: str) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
    def deco(func: Callable[..., Any]) -> Callable[..., Any]:
        @wraps(func)
        def wrapper(*a: Any, **k: Any) -> Any:
            print(f"[{name}] {func.__name__}")
            return func(*a, **k)
        return wrapper
    return deco

class Util:
    @staticmethod        # 外側: メソッド化
    @tag("static")       # 内側: 素の関数を装飾
    def identity(x: int) -> int:
        return x

    @classmethod         # 外側: メソッド化
    @tag("class")        # 内側: 素の関数を装飾
    def from_str(cls, s: str) -> tuple[type, str]:
        return (cls, s)

print(Util.identity(7))
print(Util.from_str("ok"))

from functools import wraps
from typing import Type, Any, Dict

def auto_repr(cls: Type[Any]) -> Type[Any]:
    """
    __repr__が未定義なら、自動でフィールドから生成する。
    PEP8に沿い、クラスはCapWords、関数はsnake_caseで命名します。
    """
    if "__repr__" in cls.__dict__:
        return cls

    def __repr__(self) -> str:
        fields = ", ".join(f"{k}={v!r}" for k, v in vars(self).items())
        return f"{cls.__name__}({fields})"

    setattr(cls, "__repr__", __repr__)
    return cls

REGISTRY: Dict[str, Type[Any]] = {}

def register(name: str):
    """クラスをレジストリに登録するクラスデコレータ"""
    def decorator(cls: Type[Any]) -> Type[Any]:
        REGISTRY[name] = cls
        return cls
    return decorator

@register("point")
@auto_repr
class Point:
    def __init__(self, x: int, y: int) -> None:
        self.x = x
        self.y = y

p = Point(2, 3)
print(p)                 # 自動生成されたrepr
print(REGISTRY["point"]) # レジストリに登録済み