Advanced·20 min·0/47 done

Python Decorators

Wrap functions with decorators and manage setup and teardown with context managers.

Why it matters: FastAPI routes, caching, retries, and torch.no_grad() all use these patterns.

Basics

Functions as Values

You can pass functions as arguments and return them from other functions.

A decorator takes a function and returns a wrapped version.

@timed above a definition means fn = timed(fn).

Tiny example

Pass a function as an argument and call it.

def apply(fn, x):
    return fn(x)

print(apply(len, "prompt"))

Basics

Write a Decorator

The wrapper usually calls the original function and adds behaviour around it.

Use @functools.wraps(fn) so the original name stays intact.

Decorators with arguments need one extra layer: a function that returns a decorator.

Tiny example

Print a message before the real call.

def noisy(fn):
    def wrap(x):
        print("calling", fn.__name__)
        return fn(x)
    return wrap

@noisy
def embed(t):
    return len(t)

print(embed("rag"))

Basics

Context Managers

with calls __enter__ on the way in and __exit__ on the way out, even after errors.

You can also write one with @contextlib.contextmanager: setup before yield, teardown after.

Real uses: open files, time a block, or disable gradients during inference.

Tiny example

Use with on a tiny custom context manager.

class Stage:
    def __enter__(self):
        print("start")
        return self
    def __exit__(self, *a):
        print("end")

with Stage():
    print("work")

Then AI / ML

AI / ML examples

Same idea, used in real AI work. Press Try in lab to run it.

AI / ML example

AI / ML: timing decorator

Measure how long a call takes without changing its body.

In AI / ML: Inference latency logs often come from timed decorators.

example.py

Python

import functools
import time

def timed(fn):
    @functools.wraps(fn)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = fn(*args, **kwargs)
        print(fn.__name__, "ms:", round((time.perf_counter() - start) * 1000, 2))
        return result
    return wrapper

@timed
def embed(texts):
    return [len(t) for t in texts]

print(embed(["python", "genai"]))

AI / ML example 2

AI / ML: cache repeated embeds

The second call with the same text skips the print.

In AI / ML: Embedding the same prompt twice in one process is a classic cache win.

example.py

Python

import functools

@functools.lru_cache(maxsize=128)
def fake_embed(text):
    print("computing:", text)
    return sum(ord(c) for c in text) % 997

print(fake_embed("python"))
print(fake_embed("python"))
print(fake_embed.cache_info())

Takeaways

  • 1A decorator wraps a function and returns a new one.
  • 2Use functools.wraps to keep the original name.
  • 3Context managers guarantee teardown with with.
Certificate