Decorators: timing, retry, context managers

import functools, time

def timed(fn):
    @functools.wraps(fn)            # keep name and docstring
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = fn(*args, **kwargs)
        took = time.perf_counter() - start
        print(f"{fn.__name__} took {took:.1f}s")
        return result
    return wrapper

@timed
def load_batch(rows):
    "Pretend to write rows to the warehouse."
    time.sleep(0.2)
    return len(rows)

print(load_batch([1, 2, 3]), load_batch.__name__)
print(load_batch.__doc__)

load_batch took 0.2s
3 load_batch
Pretend to write rows to the warehouse.

import functools, time

def retry(times=3, delay=0.01, on=(ConnectionError,)):
    def decorator(fn):
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            for attempt in range(1, times + 1):
                try:
                    return fn(*args, **kwargs)
                except on as e:
                    print(f"attempt {attempt} failed: {e}")
                    if attempt == times:
                        raise
                    time.sleep(delay * 2 ** (attempt - 1))
        return wrapper
    return decorator

calls = {"n": 0}
@retry(times=3)
def fetch_page():
    calls["n"] += 1
    if calls["n"] < 3:
        raise ConnectionError("warehouse timeout")
    return "page ok"
print(fetch_page())

attempt 1 failed: warehouse timeout
attempt 2 failed: warehouse timeout
page ok

from contextlib import contextmanager

@contextmanager
def step(name):
    print(f"start {name}")
    try:
        yield
        print(f"done  {name}")
    except Exception as e:
        print(f"fail  {name}: {e}")
        raise
    finally:
        print(f"cleanup {name}")
with step("extract"):
    rows = [1, 2, 3]
try:
    with step("load"):
        raise TimeoutError("warehouse timeout")
except TimeoutError:
    pass

start extract
done  extract
cleanup extract
start load
fail  load: warehouse timeout
cleanup load