# Building Context Managers: enter, exit, and @contextmanager

* * *

In [Part 2,](https://blog.moussaamzat.dev/python-with-statement-and-open-in-depth) we used `with open(...)` as a black box: it works, the file closes, everyone's happy. In [Part 3](https://blog.moussaamzat.dev/python-generator-functions-yield-lazy-evaluation), we learned how generators pause and resume execution with `yield`. Now we put both together and open the black box.

This article is about building your own context managers from scratch, understanding exactly what Python calls and when, and choosing between two approaches depending on what you're building.

* * *

## What Python Actually Does with `with`

When Python encounters a `with` statement, it doesn't do anything magical. It calls two methods in a specific order:

```python
with open("notes.txt", "r") as f:
    content = f.read()
```

Python translates this into roughly:

```python
manager = open("notes.txt", "r")
f = manager.__enter__()          # setup
try:
    content = f.read()
finally:
    manager.__exit__(None, None, None)  # teardown — always runs
```

That's the entire contract. Any object that has `__enter__` and `__exit__` methods **is** a context manager. Nothing more required.

* * *

## `__enter__`: The Setup Method

`__enter__` runs the moment execution enters the `with` block. Whatever it returns is what gets bound to the `as` name.

```python
class Example:
    def __enter__(self):
        print("entering")
        return "hello"   # this is what `as x` receives

    def __exit__(self, exc_type, exc_value, traceback):
        print("exiting")
        return False

with Example() as x:
    print(x)             # "hello"
```

Output:

```plaintext
entering
hello
exiting
```

`__enter__` can return anything: the object itself, a resource it manages, or nothing at all. For file objects, `__enter__` returns `self`, the file object, which is why `f` in `as f` is the file you write to.

* * *

## `__exit__`: The Teardown Method

`__exit__` always runs when the block exits. It receives three arguments that describe how the block exited:

```python
def __exit__(self, exc_type, exc_value, traceback):
    ...
```

| Parameter | What it contains |
| --- | --- |
| `exc_type` | The exception class, e.g. `ValueError` |
| `exc_value` | The exception instance |
| `traceback` | The traceback object |

If the block exited normally, that means no exception and all three are `None`.

**The return value matters.** Return `True` to suppress the exception. Return `False` (or `None`) to let it propagate:

```python
def __exit__(self, exc_type, exc_value, traceback):
    return True   # swallows any exception
    return False  # re-raises any exception
```

Getting this wrong is a common mistake as we'll see it again shortly.

* * *

## Building a Class-Based Context Manager

Let's build something real. A `DatabaseConnection` class that opens a SQLite connection on entry, commits on clean exit, and rolls back on error:

```python
import sqlite3

class DatabaseConnection:
    def __init__(self, db_path):
        self.db_path = db_path
        self.connection = None

    def __enter__(self):
        self.connection = sqlite3.connect(self.db_path)
        return self.connection           # caller gets the connection

    def __exit__(self, exc_type, exc_value, traceback):
        if exc_type is not None:
            self.connection.rollback()  # something went wrong — undo
        else:
            self.connection.commit()    # all good — save changes
        self.connection.close()
        return False                    # never suppress exceptions
```

Using it:

```python
with DatabaseConnection("app.db") as conn:
    cursor = conn.cursor()
    cursor.execute("CREATE TABLE IF NOT EXISTS users (name TEXT, age INT)")
    cursor.execute("INSERT INTO users VALUES ('Moussa', 30)")
# commit and close happen automatically

# If an exception occurred inside the block:
try:
    with DatabaseConnection("app.db") as conn:
        cursor = conn.cursor()
        cursor.execute("INSERT INTO users VALUES ('Moussa', 30)")
        raise ValueError("something went wrong")
except ValueError:
    pass
# rollback and close happened automatically — database unchanged
```

This pattern — commit on success, rollback on failure, always close — is exactly how Django's database transaction handling works internally.

* * *

## A Common Mistake: The Wrong Return Value

Here's a bug that's easy to miss:

```python
import time

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

    def __exit__(self, exc_type, exc_value, traceback):
        self.elapsed = time.time() - self.start
        return self.elapsed    # ⚠️ this is a float — truthy!
```

`self.elapsed` is something like `0.234` — a truthy value. So `__exit__` suppresses every exception that occurs inside the `with` block:

```python
with Timer() as t:
    raise ValueError("something broke")

print("this prints — exception silently swallowed")  # 😬
```

Always be explicit:

```python
def __exit__(self, exc_type, exc_value, traceback):
    self.elapsed = time.time() - self.start
    return False   # never suppress
```

<div data-node-type="callout">
<div data-node-type="callout-emoji">💡</div>
<div data-node-type="callout-text">Please run the code to better understand it. It is crucial for the rest of the article.</div>
</div>

* * *

## The `@contextmanager` Approach

The class approach works well, but it has a structural awkwardness: the setup and teardown are split across two methods. You have to store intermediate state, like `self.start` in the Timer, as instance variables just to pass it between `__enter__` and `__exit__`.

`@contextmanager` from `contextlib` solves this by letting you write a context manager as a single generator function. The `yield` is the dividing line:

```python
from contextlib import contextmanager

@contextmanager
def timer():
    start = time.time()      # setup — before yield
    yield                    # with block runs here
    elapsed = time.time() - start  # teardown — after yield
    print(f"Elapsed: {elapsed:.4f}s")
```

The setup and teardown live in the same function, share the same local scope, and flow naturally top to bottom. No class needed, no instance variables.

* * *

## How `@contextmanager` Works Internally

This is where [Article 3](https://blog.moussaamzat.dev/python-generator-functions-yield-lazy-evaluation) pays off. When Python sees `yield` inside a function, that function becomes a generator. `@contextmanager` wraps that generator and wires it to `__enter__` and `__exit__`:

```plaintext
__enter__ is called
    → runs the generator up to yield
    → pauses there
    → returns the yielded value (bound to `as`)

your with block runs

__exit__ is called
    → resumes the generator after yield
    → runs the rest of the function (cleanup)
```

Simplified, this is what `@contextmanager` does internally:

```python
class _GeneratorContextManager:
    def __init__(self, gen):
        self.gen = gen

    def __enter__(self):
        return next(self.gen)    # run up to yield, return yielded value

    def __exit__(self, exc_type, exc_value, traceback):
        try:
            next(self.gen)       # resume after yield — run cleanup
        except StopIteration:
            pass                 # generator finished normally
        return False
```

It's just a generator being driven by two `next()` calls. The `yield` really is a pause button.

* * *

## Yielding a Value

Whatever you `yield` is what the caller receives via `as`:

```python
@contextmanager
def database_connection(db_path):
    conn = sqlite3.connect(db_path)
    try:
        yield conn               # caller gets the connection
        conn.commit()
    except Exception:
        conn.rollback()
        raise
    finally:
        conn.close()

with database_connection("app.db") as conn:
    cursor = conn.cursor()
    cursor.execute("INSERT INTO users VALUES ('Moussa', 30)")
```

Compare this to the class version from earlier. Same behavior, fewer lines, and the logic reads top to bottom instead of being split across two methods.

* * *

## Handling Exceptions with `@contextmanager`

When an exception occurs inside the `with` block, it gets **thrown into the generator at the yield point**. Wrap the yield in `try/except` to handle it:

```python
@contextmanager
def managed_operation(name):
    print(f"Starting {name}")
    try:
        yield
    except ValueError as e:
        print(f"Handled ValueError: {e}")
        # not re-raising — exception suppressed
    except Exception:
        print("Unhandled exception — re-raising")
        raise             # re-raise anything else
    finally:
        print(f"Cleaning up {name}")  # always runs


with managed_operation("task A"):
    print("Working...")

print("---")

with managed_operation("task B"):
    raise ValueError("something went wrong")

print("Execution continues — ValueError was suppressed")
```

Output:

```plaintext
Starting task A
Working...
Cleaning up task A
---
Starting task B
Handled ValueError: something went wrong
Cleaning up task B
Execution continues — ValueError was suppressed
```

The `finally` block is your guarantee that cleanup runs regardless of what happened. Use `except` only when you specifically want to handle or suppress a particular exception.

* * *

## The One Rule: Yield Exactly Once

`@contextmanager` requires the generator to `yield` exactly one time. Not zero times. Not two times. Once.

```python
# ❌ no yield — __enter__ will crash
@contextmanager
def broken():
    print("setup")
    # forgot yield

# ❌ two yields — __exit__ will crash
@contextmanager
def also_broken():
    yield "first"
    yield "second"   # not allowed

# ✅ exactly one yield
@contextmanager
def correct():
    print("setup")
    yield
    print("teardown")
```

If you `yield` zero times, `__enter__` raises `RuntimeError`. If you `yield` more than once, `__exit__` raises `RuntimeError`. The decorator enforces this strictly.

* * *

## When to Use Each Approach

Both approaches produce objects that work identically with `with`. The choice is about which one fits the situation better.

| Situation | Use |
| --- | --- |
| Simple setup/teardown, no extra state | `@contextmanager` |
| Complex logic with multiple methods | Class with `__enter__`/`__exit__` |
| Need to subclass or extend the manager | Class approach |
| Sharing state between enter and exit | `@contextmanager` (use local variables) |
| Most real-world cases | `@contextmanager` |

In practice, `@contextmanager` covers the vast majority of use cases. The class approach is mainly useful when the context manager is complex enough to warrant its own type, or when you need to build a hierarchy of context managers.

* * *

## A Quick Side-by-Side

The same context manager written both ways:

```python
# Class-based
class ManagedFile:
    def __init__(self, path):
        self.path = path

    def __enter__(self):
        self.file = open(self.path, "w", encoding="utf-8")
        return self.file

    def __exit__(self, exc_type, exc_value, traceback):
        self.file.close()
        return False


# Function-based
from contextlib import contextmanager

@contextmanager
def managed_file(path):
    file = open(path, "w", encoding="utf-8")
    try:
        yield file
    finally:
        file.close()
```

Both work identically. The function-based version is 5 lines shorter and reads more naturally. Unless you have a specific reason to use a class, `@contextmanager` is the default choice.

* * *

## The Mental Model to Take Away

The `with` statement calls `__enter__` on the way in and `__exit__` on the way out — always. That's the entire contract.

`@contextmanager` wires a generator function to that contract automatically. Everything before `yield` maps to `__enter__`. Everything after `yield` maps to `__exit__`. The generator's pause mechanism is exactly what makes this work, and now you understand why.

Three things to carry forward:

*   `__exit__` receives exception info — check `exc_type is not None` to detect errors
    
*   Return `False` from `__exit__` unless you specifically intend to suppress an exception
    
*   With `@contextmanager`, always wrap `yield` in `try/finally` when there's cleanup to do
    

* * *

## Acronyms Used in This Article

*   **PEP** — Python Enhancement Proposal. [PEP 343](https://peps.python.org/pep-0343/) introduced the `with` statement and the context manager protocol.
    

* * *

> *This is Part 4 of 5 in the Python Context Managers series.* *Next up:* ***Part 5 — Context Managers in the Wild***