General-Purpose Context Manager Recipes

with open("prompts.jsonl", encoding="utf-8") as f:          # always pass encoding for text
    prompts = [json.loads(line) for line in f]

from pathlib import Path
text = Path("system_prompt.txt").read_text(encoding="utf-8")  # opens + closes for you

with open("data.bin", "rb") as src, open("data.bak", "wb") as dst:
    shutil.copyfileobj(src, dst)

import gzip, tarfile, zipfile

with gzip.open("corpus.jsonl.gz", "rt", encoding="utf-8") as f:
    for line in f:
        ...

with zipfile.ZipFile("dataset.zip") as zf, zf.open("train.jsonl") as member:
    first = member.readline()

with tarfile.open("model.tar.gz", "r:gz") as tf:
    tf.extractall("weights/", filter="data")   # filter= on 3.12+ (and backports); blocks path traversal

import os
import tempfile

# Hand a real file path to a library that insists on a filename (e.g. an audio
# transcription SDK), cross-platform:
with tempfile.TemporaryDirectory() as tmp:
    path = os.path.join(tmp, "clip.wav")
    with open(path, "wb") as f:
        f.write(audio_bytes)
    text = transcribe(path)

# 3.12+: NamedTemporaryFile that can be reopened by name on Windows too
with tempfile.NamedTemporaryFile("wb", suffix=".wav", delete_on_close=False) as f:
    f.write(audio_bytes)
    f.close()
    text = transcribe(f.name)

import os
from contextlib import contextmanager
from typing import Iterator


@contextmanager
def working_directory(path: str) -> Iterator[str]:
    previous = os.getcwd()
    os.chdir(path)
    try:
        yield path
    finally:
        os.chdir(previous)

import threading

_index_lock = threading.Lock()          # module-level: shared by every caller


def add_to_index(index, vectors, ids) -> None:
    with _index_lock:                   # acquire; release on any exit
        index.add_with_ids(vectors, ids)

if not _index_lock.acquire(timeout=5):
    raise TimeoutError("index busy")
try:
    index.add(vectors)
finally:
    _index_lock.release()

from contextlib import contextmanager
from typing import Iterator


@contextmanager
def acquired(lock: threading.Lock, timeout: float) -> Iterator[None]:
    if not lock.acquire(timeout=timeout):
        raise TimeoutError(f"could not acquire lock within {timeout}s")
    try:
        yield
    finally:
        lock.release()

from concurrent.futures import ThreadPoolExecutor, as_completed

with ThreadPoolExecutor(max_workers=16) as pool:        # exit = shutdown(wait=True)
    futures = {pool.submit(fetch_page, url): url for url in urls}
    for fut in as_completed(futures):
        pages[futures[fut]] = fut.result()

from multiprocessing import Pool

with Pool(8) as pool:
    chunks = pool.map(chunk_document, documents)   # map blocks until done: safe

import subprocess

result = subprocess.run(["git", "status"], capture_output=True, text=True, timeout=30, check=True)

with subprocess.Popen(
    ["python", "-u", "train.py"], stdout=subprocess.PIPE, text=True, bufsize=1,
) as proc:
    try:
        for line in proc.stdout:
            log_metric_line(line)
    except BaseException:
        proc.kill()                 # don't wait forever for a child we are abandoning
        raise
if proc.returncode:
    raise subprocess.CalledProcessError(proc.returncode, proc.args)

import socket

with socket.create_connection(("localhost", 6379), timeout=2) as sock:
    sock.sendall(b"PING\r\n")
    reply = sock.recv(64)

import requests

with requests.Session() as s:                     # reuse TCP/TLS connections
    s.headers["Authorization"] = f"Bearer {token}"
    pages = [s.get(u, timeout=10).json() for u in urls]

with requests.get(url, stream=True, timeout=30) as r:   # streamed body must be released
    r.raise_for_status()
    with open(dest, "wb") as f:
        for block in r.iter_content(1 << 20):
            f.write(block)

from urllib.request import urlopen
with urlopen("https://example.com/model-card.json", timeout=10) as resp:
    card = json.load(resp)

from unittest import mock

with mock.patch("app.llm.client.complete", return_value="stubbed answer") as fake:
    assert answer("hi") == "stubbed answer"
    fake.assert_called_once()

with mock.patch.object(settings, "max_tokens", 16):
    ...

with mock.patch.dict(PROMPTS, {"system": "You are terse."}):
    ...

import logging
from contextlib import contextmanager
from typing import Iterator, List


@contextmanager
def log_level(logger_name: str, level: int) -> Iterator[logging.Logger]:
    logger = logging.getLogger(logger_name)
    old = logger.level
    logger.setLevel(level)
    try:
        yield logger
    finally:
        logger.setLevel(old)


@contextmanager
def captured_logs(logger_name: str = "", level: int = logging.INFO) -> Iterator[List[logging.LogRecord]]:
    records: List[logging.LogRecord] = []

    class _ListHandler(logging.Handler):
        def emit(self, record: logging.LogRecord) -> None:
            records.append(record)

    handler = _ListHandler(level)
    logger = logging.getLogger(logger_name)
    logger.addHandler(handler)
    try:
        yield records
    finally:
        logger.removeHandler(handler)


with log_level("httpx", logging.WARNING):          # silence per-request INFO lines
    run_batch()

import warnings
with warnings.catch_warnings():                       # restores filters on exit (not thread-safe)
    warnings.simplefilter("ignore", category=FutureWarning)
    tokenizer = load_tokenizer()

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    legacy_call()
assert any(issubclass(w.category, DeprecationWarning) for w in caught)

from decimal import Decimal, localcontext
with localcontext() as ctx:                           # thread-local, safe
    ctx.prec = 6
    cost = Decimal(tokens) * Decimal("0.0000006")

import numpy as np
with np.errstate(divide="ignore", invalid="ignore"):
    cosine = dot / (norm_a * norm_b)                  # zero vectors -> nan, no warning

import signal
from contextlib import contextmanager
from typing import Iterator


@contextmanager
def time_limit(seconds: float) -> Iterator[None]:
    def _raise(signum, frame):
        raise TimeoutError(f"timed out after {seconds}s")

    previous = signal.signal(signal.SIGALRM, _raise)
    signal.setitimer(signal.ITIMER_REAL, seconds)
    try:
        yield
    finally:
        signal.setitimer(signal.ITIMER_REAL, 0)
        signal.signal(signal.SIGALRM, previous)

import time
from contextlib import contextmanager
from typing import Dict, Iterator


@contextmanager
def timed(label: str, sink: Dict[str, float]) -> Iterator[None]:
    t0 = time.perf_counter()
    try:
        yield
    finally:
        sink[label] = sink.get(label, 0.0) + time.perf_counter() - t0


timings: Dict[str, float] = {}
with timed("chunk", timings):
    chunks = chunker.split(doc)
with timed("embed", timings):
    vecs = embedder.encode(chunks)

import queue
from contextlib import contextmanager
from typing import Callable, Generic, Iterator, TypeVar

T = TypeVar("T")


class Pool(Generic[T]):
    """Fixed-size pool of expensive objects (e.g. tokenizers, model replicas)."""

    def __init__(self, factory: Callable[[], T], size: int) -> None:
        self._items: "queue.Queue[T]" = queue.Queue()
        for _ in range(size):
            self._items.put(factory())

    @contextmanager
    def checkout(self, timeout: float = 30.0) -> Iterator[T]:
        item = self._items.get(timeout=timeout)     # raises queue.Empty on timeout
        try:
            yield item
        finally:
            self._items.put(item)                   # always returned, even on error


tokenizers = Pool(lambda: load_tokenizer("bpe.json"), size=4)
with tokenizers.checkout() as tok:
    ids = tok.encode(text)

from contextlib import contextmanager
from typing import Iterator


class TreePrinter:
    def __init__(self) -> None:
        self.level = 0

    @contextmanager
    def indented(self) -> Iterator[None]:
        self.level += 1
        try:
            yield
        finally:
            self.level -= 1

    def line(self, text: str) -> None:
        print("  " * self.level + text)


p = TreePrinter()
p.line("agent")
with p.indented():
    p.line("tool: search")
    with p.indented():
        p.line("3 results")
p.line("done")