The Context Manager Protocol

manager = EXPR
enter = type(manager).__enter__      # looked up on the TYPE, not the instance
exit_ = type(manager).__exit__
value = enter(manager)
hit_except = False
try:
    TARGET = value                   # assignment happens inside the try
    BODY
except BaseException as exc:
    hit_except = True
    if not exit_(manager, type(exc), exc, exc.__traceback__):
        raise                        # re-raise unless __exit__ returned truthy
finally:
    if not hit_except:
        exit_(manager, None, None, None)

class RagSession:
    """Opens a vector-store connection and an LLM HTTP client together."""

    def __init__(self, dsn: str, base_url: str) -> None:
        self.dsn = dsn
        self.base_url = base_url

    def __enter__(self) -> "RagSession":
        self.db = connect_vector_store(self.dsn)
        try:
            self.http = httpx.Client(base_url=self.base_url, timeout=30)
        except BaseException:
            self.db.close()          # __exit__ will NOT run if __enter__ raises
            raise
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        try:
            self.http.close()
        finally:
            self.db.close()          # close both even if the first close fails

from contextlib import ExitStack


class RagSession:
    def __init__(self, dsn: str, base_url: str) -> None:
        self.dsn, self.base_url = dsn, base_url
        self._stack = ExitStack()

    def __enter__(self) -> "RagSession":
        with ExitStack() as stack:                       # unwinds if anything below fails
            self.db = stack.enter_context(closing(connect_vector_store(self.dsn)))
            self.http = stack.enter_context(httpx.Client(base_url=self.base_url))
            self._stack = stack.pop_all()                # success: transfer ownership
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self._stack.close()

class PromptCache:
    def __enter__(self) -> "PromptCache":     # annotate with the class (string or `Self`)
        self._conn = sqlite3.connect(self.path)
        return self

from types import TracebackType
from typing import Optional, Type


class SkipOnRateLimit:
    """Swallow RateLimitError only; count how many calls were skipped."""

    def __init__(self) -> None:
        self.skipped = 0

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

    def __exit__(
        self,
        exc_type: Optional[Type[BaseException]],
        exc: Optional[BaseException],
        tb: Optional[TracebackType],
    ) -> bool:
        if exc_type is not None and issubclass(exc_type, RateLimitError):
            self.skipped += 1
            return True          # handled: suppress
        return False             # everything else propagates


skipper = SkipOnRateLimit()
for doc in docs:
    with skipper:
        summaries.append(llm.summarize(doc))
print(f"skipped {skipper.skipped} docs due to rate limits")

class StepContext:
    def __init__(self, step: str) -> None:
        self.step = step

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

    def __exit__(self, exc_type, exc, tb) -> None:
        if exc is not None:
            exc.add_note(f"while running pipeline step {self.step!r}")   # 3.11+
        # returning None: the (now annotated) exception propagates

class Upload:
    def __exit__(self, exc_type, exc, tb) -> None:
        self.conn.close()              # if this raises, the body's error is now only __context__

import logging

log = logging.getLogger(__name__)


class Checkpointer:
    def __exit__(self, exc_type, exc, tb) -> None:
        try:
            self.flush()
        except Exception:
            if exc is None:
                raise                      # nothing else went wrong: surface the flush error
            log.exception("flush failed while handling %r", exc)   # keep the original error primary
        finally:
            self.file.close()

with open(src) as fin, open(dst, "w") as fout:       # any version
    fout.write(fin.read())

with (                                                # 3.10+ official (3.9 parser accepts it unofficially)
    open("prompts.jsonl") as prompts,
    open("completions.jsonl", "w") as out,
    torch.inference_mode(),
):
    ...

lock = threading.Lock()
with lock:
    with lock:        # deadlock: Lock is reusable but NOT re-entrant
        ...

rlock = threading.RLock()
with rlock:
    with rlock:       # fine
        ...

import threading


class ModelLease:
    """Nested `with lease:` blocks share one GPU load; unload when the outermost exits."""

    def __init__(self, loader) -> None:
        self._loader = loader
        self._depth = 0
        self._lock = threading.RLock()
        self.model = None

    def __enter__(self):
        with self._lock:
            if self._depth == 0:
                self.model = self._loader()
            self._depth += 1
            return self.model

    def __exit__(self, exc_type, exc, tb) -> None:
        with self._lock:
            self._depth -= 1
            if self._depth == 0:
                self.model = None          # drop the reference so GPU memory can be freed

class OneShot:
    def __init__(self) -> None:
        self._used = False

    def __enter__(self) -> "OneShot":
        if self._used:
            raise RuntimeError("OneShot can only be entered once; create a new instance")
        self._used = True
        return self

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

from __future__ import annotations

from contextlib import AbstractContextManager, contextmanager
from typing import Iterator


@contextmanager
def batch_writer(path: str, batch_size: int = 512) -> Iterator[BatchWriter]:
    writer = BatchWriter(path, batch_size)
    try:
        yield writer
    finally:
        writer.flush_and_close()


def run_eval(guard: AbstractContextManager[object]) -> None:
    with guard:
        ...

class VectorStore:
    @contextmanager
    def transaction(self) -> Iterator["VectorStore"]:
        self._conn.execute("BEGIN")
        try:
            yield self
        except BaseException:
            self._conn.rollback()
            raise
        else:
            self._conn.commit()