Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""A lightweight implementation of `@contextlib.contextmanager`."""
3from __future__ import annotations
5from collections.abc import Callable, Generator, Iterator
6from functools import wraps
7from types import TracebackType
8from typing import Generic, ParamSpec, TypeVar
10_T = TypeVar('_T')
11_P = ParamSpec('_P')
14class _GeneratorContextManager(Generic[_T]):
15 """A lighter version of `contextlib._GeneratorContextManager`.
17 Unlike the standard library implementation, it can't be used as a decorator or be recreated.
18 This avoids storing the generator function and its arguments on each call.
19 """
21 __slots__ = ('gen',)
23 def __init__(self, gen: Generator[_T]) -> None:
24 self.gen = gen
26 def __enter__(self) -> _T:
27 # Faster than `next(self.gen)` (see https://github.com/python/cpython/pull/141275):
28 for value in self.gen:
29 return value
30 raise RuntimeError("generator didn't yield") # pragma: no cover
32 def __exit__(
33 self, typ: type[BaseException] | None, value: BaseException | None, traceback: TracebackType | None
34 ) -> bool | None: # Same return type as `contextlib`, so that type checkers don't assume exceptions are suppressed
35 if typ is None:
36 # Faster than calling `next(self.gen)` and checking for `StopIteration`:
37 for _ in self.gen:
38 try: # pragma: no cover
39 raise RuntimeError("generator didn't stop")
40 finally:
41 self.gen.close() # pragma: no cover
42 return False
44 # Same logic as `contextlib._GeneratorContextManager.__exit__()`:
45 if value is None:
46 # Need to force instantiation so we can reliably tell if we get the same exception back:
47 value = typ() # pragma: no cover
48 try:
49 self.gen.throw(value)
50 except StopIteration as exc:
51 # Suppress `StopIteration` *unless* it's the same exception that was passed to `throw()`:
52 return exc is not value # pragma: no cover
53 except RuntimeError as exc:
54 # Don't re-raise the passed in exception:
55 if exc is value:
56 exc.__traceback__ = traceback
57 return False
58 # Avoid suppressing if a `StopIteration` exception was passed to `throw()` and later
59 # wrapped into a `RuntimeError` (see PEP 479 for sync generators):
60 if isinstance(value, StopIteration) and exc.__cause__ is value: # pragma: no cover
61 value.__traceback__ = traceback
62 return False
63 raise # pragma: no cover
64 except BaseException as exc:
65 # Only re-raise if it's *not* the exception that was passed to `throw()`, because
66 # `__exit__()` must not raise an exception unless `__exit__()` itself failed:
67 if exc is not value:
68 raise # pragma: no cover
69 exc.__traceback__ = traceback
70 return False
71 try: # pragma: no cover
72 raise RuntimeError("generator didn't stop after throw()")
73 finally:
74 self.gen.close() # pragma: no cover
77def contextmanager(func: Callable[_P, Iterator[_T]]) -> Callable[_P, _GeneratorContextManager[_T]]:
78 """A lightweight replacement for `contextlib.contextmanager`, only usable in `with` statements."""
80 @wraps(func)
81 def helper(*args: _P.args, **kwargs: _P.kwargs) -> _GeneratorContextManager[_T]:
82 return _GeneratorContextManager(func(*args, **kwargs)) # pyright: ignore[reportArgumentType]
84 return helper