# Decorators & Context Managers: The @ and with You Keep Seeing

Open almost any modern Python project and you'll spot two things quickly. Lines that start with `@`, like `@app.get("/hello")`. And blocks that start with `with`, like `with torch.no_grad():`.

They look like magic if nobody explains them. They're not. Both are built from ideas you already met: functions as values and closures from Post 6, and generators from Post 7. This post is core, because FastAPI, PyTorch, and plenty of AI libraries lean on both.

## The problem decorators solve

Say you want to know how long a function takes. You could add timing code to it. Then you want the same for five other functions. Then you want to add logging. Soon every function is half timing code and half the thing it was meant to do.

What you'd like is a way to wrap extra behavior around a function _without touching its insides_. That's what a decorator does.

## Building a decorator by hand

A decorator is a function that takes a function and returns a new, improved function. Here's one that measures time:

```python
import time

def timed(func):
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        elapsed = time.perf_counter() - start
        print(f"{func.__name__} took {elapsed:.3f} seconds")
        return result
    return wrapper
```

Take it slowly, because everything in there comes from Post 6:

-   `timed` takes a function called `func`.
    
-   Inside, it defines `wrapper`, a new function that starts a clock, calls the original `func`, stops the clock, prints the time, and returns the original result.
    
-   `*args` and `**kwargs` let the wrapper accept whatever the original function accepts and pass it along.
    
-   `wrapper` is a closure. It remembers `func` after `timed` has finished.
    
-   `timed` returns `wrapper`, without calling it.
    

Now use it:

```python
def slow_add(a, b):
    time.sleep(0.5)
    return a + b

slow_add = timed(slow_add)

print(slow_add(2, 3))
# slow_add took 0.500 seconds
# 5
```

We replaced `slow_add` with the wrapped version. Python has a shortcut for exactly this pattern, and that's the `@` symbol:

```python
@timed
def slow_add(a, b):
    time.sleep(0.5)
    return a + b
```

Those two versions do the same thing. `@timed` above a function simply means `slow_add = timed(slow_add)`. That's all it is.

## A small fix: functools.wraps

There's one annoying side effect. After decorating, the function forgets who it is:

```python
print(slow_add.__name__)    # wrapper
```

That's because `slow_add` is now really the `wrapper` function. It can mess up error messages, logs, and tools that look at function names. The standard fix is to add `@wraps(func)` to the wrapper:

```python
from functools import wraps

def timed(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.perf_counter()
        result = func(*args, **kwargs)
        elapsed = time.perf_counter() - start
        print(f"{func.__name__} took {elapsed:.3f} seconds")
        return result
    return wrapper
```

Now `slow_add.__name__` correctly says `slow_add`. Treat `@wraps` as a habit you always include. It's a decorator that's used inside a decorator, which is a fun thing to notice.

## Decorators that take settings

Often you want to customize a decorator, like `@retry(attempts=5)`. That needs one more layer, and this is the part that confuses everyone at first. Here's the retry helper from Post 6, rewritten as a decorator:

```python
def retry(attempts=3, delay=0.5):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(1, attempts + 1):
                try:
                    return func(*args, **kwargs)
                except Exception as error:
                    print(f"Attempt {attempt} failed: {error}")
                    if attempt == attempts:
                        raise
                    time.sleep(delay)
        return wrapper
    return decorator
```

Three layers, each with one job:

1.  `retry(...)` takes the settings and returns a decorator.
    
2.  `decorator(func)` takes the function and returns the wrapper.
    
3.  `wrapper(...)` is what actually runs each time you call the function.
    

So `@retry(attempts=5)` first calls `retry(attempts=5)`, which hands back a decorator, and then that decorator is applied to the function below it. If three layers feels like a lot, it is. Read it twice, then type it out yourself.

## Decorators that register instead of wrap

Not every decorator changes the function. Some just make a note of it and hand it back unchanged. This is how web frameworks find your route handlers, and it's how many AI frameworks keep track of the tools a model is allowed to call.

```python
tools = {}

def tool(func):
    tools[func.__name__] = func
    return func

@tool
def get_weather(city):
    return f"Sunny in {city}"

print(tools["get_weather"]("Lisbon"))    # Sunny in Lisbon
print(list(tools))                       # ['get_weather']
```

Remember the dictionary of functions from Post 6? This is the tidy version. Mark a function with `@tool` and it lands in the registry. When a model asks for `get_weather` by name, you look it up and run it.

When you see `@app.get("/hello")` in FastAPI, it's the same family of idea: it tells the framework "when someone visits `/hello`, run this function."

## A few built-in decorators worth knowing

-   `@cache` (from `functools`) remembers results, so calling a function again with the same input returns instantly:
    

```python
from functools import cache

@cache
def expensive_lookup(word):
    print(f"Computing for {word}...")
    return len(word) * 10

print(expensive_lookup("hello"))    # Computing for hello...  then 50
print(expensive_lookup("hello"))    # 50  (no computing this time)
```

Two warnings. The inputs must be things that can be dictionary keys, so passing a list fails (remember Post 5). And don't cache anything where you _want_ a fresh answer each time, like most model calls.

-   `@staticmethod`, `@property`, and `@dataclass` show up when we cover classes in the next post.
    

## Context managers: cleaning up after yourself

Now the `with` statement. Some things need to be opened and then closed properly: files, network connections, database sessions, locks. If you forget to close them, or your program crashes in the middle, you get leaks and strange bugs.

You've already seen the fix:

```python
with open("demo.txt", "w", encoding="utf-8") as file:
    file.write("hello")

print(file.closed)    # True
```

When the `with` block ends, the file is closed. Always. Even if an error happens inside the block. That guarantee is the whole point.

You can open more than one thing in the same line:

```python
with open("demo.txt") as src, open("copy.txt", "w") as dst:
    dst.write(src.read())
```

## How `with` works

Anything that works with `with` is called a context manager. It's an object with two special methods: one that runs when you enter the block, and one that runs when you leave, no matter how you leave. Here's a timer built that way:

```python
import time

class Timer:
    def __enter__(self):
        self.start = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self.elapsed = time.perf_counter() - self.start
        print(f"Took {self.elapsed:.2f} seconds")
        return False

with Timer():
    time.sleep(0.3)
# Took 0.30 seconds
```

We haven't covered classes yet (that's Post 9), so just read this as a recipe for now. `__enter__` is the setup. `__exit__` is the cleanup. If an error happened inside the block, the three extra arguments describe it, and returning `False` means "let the error carry on."

## The easy way: `@contextmanager`

Writing a whole class for a small job is a bit much. Python has a shortcut that uses a generator, which you met in Post 7. Everything before `yield` is setup, and everything after is cleanup:

```python
from contextlib import contextmanager
import time

@contextmanager
def timer(label):
    start = time.perf_counter()
    try:
        yield
    finally:
        elapsed = time.perf_counter() - start
        print(f"{label} took {elapsed:.2f} seconds")

with timer("sleeping"):
    time.sleep(0.3)
# sleeping took 0.30 seconds
```

The `try` and `finally` matter. Code in `finally` always runs, even if the block inside `with` raises an error. So the timer still reports in a crash:

```python
with timer("crashing"):
    raise ValueError("oops")
# crashing took 0.00 seconds
# then the ValueError appears as usual
```

Notice how this ties together. `@contextmanager` is a decorator, wrapped around a generator, producing something you use with `with`. If that sentence made sense, you've understood the last three posts.

## Where you'll see these in AI work

-   `with torch.no_grad():` switches off gradient tracking while you run a model for predictions, then switches it back on afterward. You'll write it often once we reach PyTorch. (It can also be used as a decorator, so you'll see both styles.)
    
-   **Network and database sessions** that must close cleanly.
    
-   **Temporary settings**, like running a block of code with a different configuration and then restoring the original.
    
-   **Timing and tracing** blocks of code, like our `timer`.
    
-   **Route handlers and tools** registered with decorators in web and agent frameworks.
    

## Mini project: a small toolkit

Let's combine everything in one file. Create `toolkit.py` with the `timed` decorator, the `retry` decorator, and the `timer` context manager from above, then add this at the bottom:

```python
import random

@timed
@retry(attempts=4, delay=0.1)
def fake_model_call(prompt):
    time.sleep(0.2)
    if random.random() < 0.5:
        raise RuntimeError("Model busy")
    return f"Answer to: {prompt}"

with timer("whole batch"):
    for prompt in ["What is a decorator?", "What is a context manager?"]:
        print(fake_model_call(prompt))
```

Run it a few times. The failures are random, so every run looks a little different. A run with one failed attempt might print:

```yaml
Attempt 1 failed: Model busy
fake_model_call took 0.500 seconds
Answer to: What is a decorator?
fake_model_call took 0.200 seconds
Answer to: What is a context manager?
whole batch took 0.70 seconds
```

A detail worth noticing: when you stack decorators, the one closest to the function is applied first. So `@retry` wraps the function, and then `@timed` wraps _that_. That means the time you see includes the retries and the waiting between them. If you swapped the two lines, `timed` would report only a single attempt each time.

You've just written the core of a tool that real AI apps need: a function that measures itself, tries again when it fails, and sits inside a timed block. And none of that code is inside the function that makes the model call.

### Make it yours

1.  Add a `@log_calls` decorator that prints the function name and its arguments before each call.
    
2.  Write your own simple cache using a dictionary, then compare it with `@cache`.
    
3.  Write a context manager that temporarily changes a value in a dictionary and puts the original back afterward.
    
4.  Add a `max_total_seconds` setting to `retry` so it gives up if too much time has passed.
    

## Practice

About 20 minutes of real effort on each before asking for help.

1.  Write a `@shout` decorator that makes any function's returned text uppercase.
    
2.  Write a `@count_calls` decorator that keeps track of how many times a function has been called. (Hint: you can store the count on the wrapper itself, like `wrapper.calls = 0`.)
    
3.  Write a context manager with `@contextmanager` that prints "Entering" when the block starts and "Leaving" when it ends, even if there's an error.
    
4.  Use `@cache` on a recursive Fibonacci function and compare how long `fib(35)` takes with and without it.
    
5.  Without using `@`, write the exact equivalent of decorating a function with `@timed`.
    

Then ask an AI tutor to review your code, not rewrite it.

## What to take away

-   A decorator is a function that takes a function and returns a new one. `@name` is just a shortcut for `func = name(func)`.
    
-   Use `@wraps(func)` inside every decorator you write.
    
-   A decorator with settings has three layers: settings, function, wrapper.
    
-   Some decorators wrap, and some simply register a function and return it unchanged.
    
-   A context manager guarantees setup and cleanup around a block of code, even when something goes wrong.
    
-   `@contextmanager` lets you build one with a generator and a `try`/`finally`.
    

* * *

**Previous:** Post 7, _Iterators, Generators & Comprehensions_

**Next in the series:** Post 9, _OOP the Practical Way_, where we cover classes, dataclasses, and when to use composition instead of inheritance.

---

*Published via [ZyVOP](https://zyvop.com/decorators-context-managers-the-and-with-you-keep-seeing-c925u?utm_source=hashnode&utm_medium=crosspost&utm_campaign=syndication) — Write once in Markdown, auto-backup to GitHub, and syndicate to Dev.to, Medium & Hashnode in 1 click.*
