contextlib Reference

from contextlib import contextmanager
from typing import Iterator


@contextmanager
def openai_timeout(client, seconds: float) -> Iterator[None]:
    """Temporarily change a client's timeout."""
    old = client.timeout
    client.timeout = seconds        # setup: before try, so a failed setup is not "undone"
    try:
        yield                       # the with-body runs here
    finally:
        client.timeout = old        # teardown: always

@contextmanager
def always_cleanup() -> Iterator[Resource]:
    r = acquire()
    try:
        yield r
    finally:
        release(r)                  # runs on success and on error


@contextmanager
def commit_or_rollback(conn) -> Iterator[Connection]:
    try:
        yield conn
    except BaseException:
        conn.rollback()
        raise                       # MUST re-raise or the error is swallowed
    else:
        conn.commit()               # only on success


@contextmanager
def log_errors(step: str) -> Iterator[None]:
    try:
        yield
    except Exception as exc:
        log.error("step %s failed: %s", step, exc)
        raise


@contextmanager
def ignore(*excs: type) -> Iterator[None]:   # deliberate suppression (= contextlib.suppress)
    try:
        yield
    except excs:
        pass

@contextmanager
def traced(name: str) -> Iterator[None]:
    start = time.perf_counter()
    try:
        yield
    finally:
        metrics.observe(name, time.perf_counter() - start)


@traced("rerank")
def rerank(query: str, docs: list) -> list:
    ...

from contextlib import asynccontextmanager
from typing import AsyncIterator

import httpx


@asynccontextmanager
async def llm_http(base_url: str) -> AsyncIterator[httpx.AsyncClient]:
    client = httpx.AsyncClient(base_url=base_url, timeout=httpx.Timeout(60, connect=5))
    try:
        yield client
    finally:
        await client.aclose()


async def main() -> None:
    async with llm_http("https://api.example.com") as http:
        r = await http.post("/v1/chat/completions", json=payload)

from contextlib import ExitStack
from pathlib import Path


def merge_shards(shard_paths: list[Path], out_path: Path) -> int:
    with ExitStack() as stack:
        shards = [stack.enter_context(p.open(encoding="utf-8")) for p in shard_paths]
        out = stack.enter_context(out_path.open("w", encoding="utf-8"))
        n = 0
        for f in shards:
            for line in f:
                out.write(line)
                n += 1
        return n
    # every file closed here, in reverse order, even if one open() failed midway

def generate(model, batch, *, use_amp: bool, profile: bool):
    with ExitStack() as stack:
        stack.enter_context(torch.inference_mode())
        if use_amp:
            stack.enter_context(torch.autocast("cuda", dtype=torch.bfloat16))
        if profile:
            prof = stack.enter_context(torch.profiler.profile())
        return model.generate(**batch)

with ExitStack() as stack:
    old_level = logger.level
    logger.setLevel(logging.DEBUG)
    stack.callback(logger.setLevel, old_level)

    tmp = Path(tempfile.mkdtemp())
    stack.callback(shutil.rmtree, tmp, ignore_errors=True)

    run_agent(workdir=tmp)

class EmbeddingService:
    def __init__(self, model_path: str, index_path: str) -> None:
        with ExitStack() as stack:
            self.model_file = stack.enter_context(open(model_path, "rb"))
            self.index = stack.enter_context(closing(load_index(index_path)))
            self._resources = stack.pop_all()     # nothing is closed on success

    def close(self) -> None:
        self._resources.close()

    def __enter__(self) -> "EmbeddingService":
        return self

    def __exit__(self, *exc_info) -> None:
        self.close()

from contextlib import AsyncExitStack


async def run_agent_with_tools(tool_configs: list[dict]) -> str:
    async with AsyncExitStack() as stack:
        http = await stack.enter_async_context(httpx.AsyncClient(timeout=30))
        sessions = [
            await stack.enter_async_context(connect_mcp_server(cfg))   # each tool server
            for cfg in tool_configs
        ]
        stack.push_async_callback(flush_traces)
        return await agent_loop(http, sessions)

from contextlib import closing
import sqlite3

with closing(sqlite3.connect("evals.db")) as conn:
    with conn:                                 # sqlite3 connection's own CM = transaction, NOT close
        conn.execute("INSERT INTO runs VALUES (?, ?)", (run_id, score))

from contextlib import aclosing


async def stream_tokens(client, prompt):
    async with client.stream("POST", "/v1/completions", json={"prompt": prompt, "stream": True}) as r:
        async for line in r.aiter_lines():
            if line.startswith("data: "):
                yield line[6:]


async def first_sentence(client, prompt) -> str:
    out = []
    async with aclosing(stream_tokens(client, prompt)) as tokens:
        async for tok in tokens:
            out.append(tok)
            if tok.endswith("."):
                break                 # aclose() runs now: the HTTP stream is closed immediately
    return "".join(out)

@asynccontextmanager
async def aclosing_compat(agen):
    try:
        yield agen
    finally:
        await agen.aclose()

from contextlib import suppress

with suppress(FileNotFoundError):
    os.remove(cache_path)

with suppress(KeyError, IndexError):
    usage = response["usage"]["completion_tokens"]

import io
from contextlib import redirect_stdout, redirect_stderr

buf = io.StringIO()
with redirect_stdout(buf), redirect_stderr(buf):
    result = run_python_tool(code)        # in-process tool call
tool_output = buf.getvalue()[-4000:]      # truncate before sending back to the LLM

from contextlib import nullcontext

def embed(texts, model, *, no_grad: bool = True):
    guard = torch.no_grad() if no_grad else nullcontext()
    with guard:
        return model(texts)


def write_results(rows, out=None):
    # caller-owned stream vs file we open and must close
    cm = open("results.jsonl", "w") if out is None else nullcontext(out)
    with cm as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")

from contextlib import chdir

with chdir(sandbox_dir):
    subprocess.run(["pytest", "-q"], check=True)

from contextlib import AbstractContextManager


class TokenBudget(AbstractContextManager):
    """__enter__ defaults to returning self; only __exit__ is abstract."""

    def __init__(self, limit: int) -> None:
        self.limit = limit
        self.used = 0

    def charge(self, n: int) -> None:
        self.used += n
        if self.used > self.limit:
            raise BudgetExceeded(self.used, self.limit)

    def __exit__(self, exc_type, exc, tb) -> None:
        metrics.gauge("tokens.used", self.used)

from contextlib import ContextDecorator


class log_latency(ContextDecorator):
    def __init__(self, name: str) -> None:
        self.name = name

    def __enter__(self) -> "log_latency":
        self.t0 = time.perf_counter()
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        log.info("%s %.1fms", self.name, (time.perf_counter() - self.t0) * 1000)


@log_latency("retrieve")
def retrieve(query: str): ...

with log_latency("rerank"):
    ...