Context Managers Explained

Introduction

Welcome back to The Python Data Model & Protocols! We've now reached our third lesson in this learning path, well done! In our previous lessons, we've seen how Python's data model provides elegant protocols for common operations. We built custom value types that behaved like built-in objects, and we created streaming data pipelines that processed information efficiently using generators. Today, we're exploring another essential protocol in Python's toolkit: context managers . Context managers solve a fundamental problem in programming: ensuring that resources are properly acquired and released, even when errors occur. Whether we're working with files, network connections, database transactions, or temporary directories, context managers guarantee that cleanup happens reliably through Python's with statement. We'll discover how to implement context managers using two distinct approaches: the traditional class-based method with__enter__and__exit__dunder methods, and the more concise generator-based approach using decorators. Our practical example will involve managing temporary working directories, demonstrating how context managers can save and restore system state while handling exceptions gracefully.

Understanding Resource Management Challenges

Before diving into context managers, let's understand why they're necessary. In any nontrivial program, we frequently work with resources that require explicit cleanup: open files need closing, network connections need termination, and temporary directories need removal. Without proper management, these resources can leak, leading to system instability or exhaustion. Traditional approaches using try-finally blocks work but become verbose and error-prone as complexity increases. We might forget to handle all error paths or accidentally suppress important exceptions during cleanup. Context managers solve these problems by providing a standardized protocol that guarantees both setup and teardown operations execute correctly. The with statement in Python automatically calls special methods on our context manager objects:__enter__runs before the code block begins, and__exit__runs after the block completes, regardless of whether an exception occurred. This pattern ensures that cleanup code always executes, making our programs more robust and predictable.

Implementing Class-Based Context Managers

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Python

from __future__ import annotations

import contextlib
import os
import tempfile
from pathlib import Path
from types import TracebackType
from typing import Iterator


class TempWorkDir:
    """
    Class-based context manager that switches to a new temporary directory
    for the duration of the with-block, then restores the previous CWD and
    removes the temporary directory. Exceptions are not suppressed.
    """

    _prev_cwd: Path
    _tmp: tempfile.TemporaryDirectory[str]

from __future__ import annotations

import contextlib
import os
import tempfile
from pathlib import Path
from types import TracebackType
from typing import Iterator


class TempWorkDir:
    """
    Class-based context manager that switches to a new temporary directory
    for the duration of the with-block, then restores the previous CWD and
    removes the temporary directory. Exceptions are not suppressed.
    """

    _prev_cwd: Path
    _tmp: tempfile.TemporaryDirectory[str]

Python

    def __enter__(self) -> Path:
        self._prev_cwd = Path.cwd()
        self._tmp = tempfile.TemporaryDirectory()
        new = Path(self._tmp.name)
        os.chdir(new)
        return new

    def __enter__(self) -> Path:
        self._prev_cwd = Path.cwd()
        self._tmp = tempfile.TemporaryDirectory()
        new = Path(self._tmp.name)
        os.chdir(new)
        return new

Python

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> bool:
        try:
            os.chdir(self._prev_cwd)
        finally:
            # Always attempt cleanup, regardless of chdir outcome or exceptions.
            self._tmp.cleanup()
        return False  # Do not suppress exceptions

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> bool:
        try:
            os.chdir(self._prev_cwd)
        finally:
            # Always attempt cleanup, regardless of chdir outcome or exceptions.
            self._tmp.cleanup()
        return False  # Do not suppress exceptions

Python

@contextlib.contextmanager
def temp_workdir() -> Iterator[Path]:
    """
    Generator-based context manager alternative that mirrors TempWorkDir.
    """
    prev = Path.cwd()
    with tempfile.TemporaryDirectory() as td:
        new = Path(td)
        try:
            os.chdir(new)
            yield new
        finally:
            os.chdir(prev)

@contextlib.contextmanager
def temp_workdir() -> Iterator[Path]:
    """
    Generator-based context manager alternative that mirrors TempWorkDir.
    """
    prev = Path.cwd()
    with tempfile.TemporaryDirectory() as td:
        new = Path(td)
        try:
            os.chdir(new)
            yield new
        finally:
            os.chdir(prev)

Python

if __name__ == "__main__":
    original = Path.cwd()

    # Demonstrate class-based manager with an intentional exception inside.
    try:
        with TempWorkDir() as d1:
            Path("a.txt").write_text("alpha")
            print(Path.cwd())
            print(Path("a.txt").is_file())
            raise RuntimeError("boom")
    except RuntimeError as e:
        print(str(e))
    print(Path.cwd() == original)
    print(Path(d1).exists())

if __name__ == "__main__":
    original = Path.cwd()

    # Demonstrate class-based manager with an intentional exception inside.
    try:
        with TempWorkDir() as d1:
            Path("a.txt").write_text("alpha")
            print(Path.cwd())
            print(Path("a.txt").is_file())
            raise RuntimeError("boom")
    except RuntimeError as e:
        print(str(e))
    print(Path.cwd() == original)
    print(Path(d1).exists())

Python

    # Demonstrate generator-based manager without exceptions.
    with temp_workdir() as d2:
        Path("b.txt").write_text("beta")
        print(Path.cwd())
        print(Path("b.txt").is_file())
    print(Path.cwd() == original)
    print(Path(d2).exists())

    # Demonstrate generator-based manager without exceptions.
    with temp_workdir() as d2:
        Path("b.txt").write_text("beta")
        print(Path.cwd())
        print(Path("b.txt").is_file())
    print(Path.cwd() == original)
    print(Path(d2).exists())

text

/tmp/tmprk0pol8l
True
boom
True
False
/tmp/tmp7q6eqkx3
True
True
False

/tmp/tmprk0pol8l
True
boom
True
False
/tmp/tmp7q6eqkx3
True
True
False

Python

from __future__ import annotations

import contextlib
import os
import tempfile
from pathlib import Path
from types import TracebackType
from typing import Iterator


class TempWorkDir:
    """
    Class-based context manager that switches to a new temporary directory
    for the duration of the with-block, then restores the previous CWD and
    removes the temporary directory. Exceptions are not suppressed.
    """

    _prev_cwd: Path
    _tmp: tempfile.TemporaryDirectory[str]

Python

    def __enter__(self) -> Path:
        self._prev_cwd = Path.cwd()
        self._tmp = tempfile.TemporaryDirectory()
        new = Path(self._tmp.name)
        os.chdir(new)
        return new

Python

    def __exit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> bool:
        try:
            os.chdir(self._prev_cwd)
        finally:
            # Always attempt cleanup, regardless of chdir outcome or exceptions.
            self._tmp.cleanup()
        return False  # Do not suppress exceptions

Python

@contextlib.contextmanager
def temp_workdir() -> Iterator[Path]:
    """
    Generator-based context manager alternative that mirrors TempWorkDir.
    """
    prev = Path.cwd()
    with tempfile.TemporaryDirectory() as td:
        new = Path(td)
        try:
            os.chdir(new)
            yield new
        finally:
            os.chdir(prev)

Python

if __name__ == "__main__":
    original = Path.cwd()

    # Demonstrate class-based manager with an intentional exception inside.
    try:
        with TempWorkDir() as d1:
            Path("a.txt").write_text("alpha")
            print(Path.cwd())
            print(Path("a.txt").is_file())
            raise RuntimeError("boom")
    except RuntimeError as e:
        print(str(e))
    print(Path.cwd() == original)
    print(Path(d1).exists())

Python

    # Demonstrate generator-based manager without exceptions.
    with temp_workdir() as d2:
        Path("b.txt").write_text("beta")
        print(Path.cwd())
        print(Path("b.txt").is_file())
    print(Path.cwd() == original)
    print(Path(d2).exists())

text

/tmp/tmprk0pol8l
True
boom
True
False
/tmp/tmp7q6eqkx3
True
True
False

Introduction

Understanding Resource Management Challenges

Implementing Class-Based Context Managers

Join the 1M+ learners on CodeSignal

The Enter Method: Setting Up Resources

The Exit Method: Guaranteed Cleanup

Generator-Based Context Managers

Demonstrating Exception Handling

Comparing Both Approaches in Action

Observing the Complete Output

Conclusion and Next Steps