Python Cheatsheet 04 — Decorators and Context Managers

from functools import wraps

def log_calls(fn):
    @wraps(fn)
    def inner(*args, **kwargs):
        print(f"calling {fn.__name__}")
        return fn(*args, **kwargs)
    return inner

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

def retry(times: int):
    def decorator(fn):
        @wraps(fn)
        def inner(*args, **kwargs):
            for i in range(times):
                try:
                    return fn(*args, **kwargs)
                except Exception:
                    if i == times - 1: raise
        return inner
    return decorator

@retry(times=3)
def flaky(): ...

def log_async(fn):
    @wraps(fn)
    async def inner(*args, **kwargs):
        print(f"async {fn.__name__}")
        return await fn(*args, **kwargs)
    return inner

@log_async
async def fetch(): ...

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

P = ParamSpec("P")
R = TypeVar("R")

def log_calls(fn: Callable[P, R]) -> Callable[P, R]:
    @wraps(fn)
    def inner(*args: P.args, **kwargs: P.kwargs) -> R:
        return fn(*args, **kwargs)
    return inner

def add_repr(cls):
    def __repr__(self):
        return f"{cls.__name__}({self.__dict__})"
    cls.__repr__ = __repr__
    return cls

@add_repr
class Point:
    def __init__(self, x, y):
        self.x, self.y = x, y

@cache
@validate_args
@log_calls
def expensive(x: int) -> int:
    return x ** 2

from functools import cache, lru_cache

@cache                    # unbounded; 3.9+
def fib(n): ...

@lru_cache(maxsize=128)
def parse(s): ...

from functools import partial

add = lambda x, y: x + y
add5 = partial(add, 5)
add5(3)                   # 8

class DBConnection:
    def __enter__(self):
        self.conn = connect()
        return self.conn
    def __exit__(self, exc_type, exc, tb):
        self.conn.close()

with DBConnection() as conn:
    conn.execute(...)

from contextlib import contextmanager

@contextmanager
def acquire(lock):
    lock.acquire()
    try:
        yield lock
    finally:
        lock.release()

with acquire(my_lock):
    ...

class AsyncDB:
    async def __aenter__(self):
        self.conn = await connect()
        return self.conn
    async def __aexit__(self, exc_type, exc, tb):
        await self.conn.close()

async with AsyncDB() as conn:
    ...

from contextlib import asynccontextmanager

@asynccontextmanager
async def lifespan(app):
    db = await connect()
    yield db
    await db.close()

from contextlib import ExitStack

with ExitStack() as stack:
    files = [stack.enter_context(open(f)) for f in paths]
    # all files closed on exit

from contextlib import AsyncExitStack

async with AsyncExitStack() as stack:
    conns = [await stack.enter_async_context(connect(d)) for d in dbs]

from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove("maybe.txt")

from contextlib import redirect_stdout
import io

buf = io.StringIO()
with redirect_stdout(buf):
    print("hello")

captured = buf.getvalue()

def singleton(cls):
    instance = None
    @wraps(cls)
    def inner(*args, **kwargs):
        nonlocal instance
        if instance is None:
            instance = cls(*args, **kwargs)
        return instance
    return inner

@singleton
class Cache:
    ...

class User:
    @property
    def full_name(self) -> str:
        return f"{self.first} {self.last}"
    
    @full_name.setter
    def full_name(self, value: str):
        self.first, self.last = value.split(" ", 1)

class Math:
    @staticmethod
    def square(x): return x ** 2
    
    @classmethod
    def from_dict(cls, d): return cls(**d)