1"""
  2The shape every `capture_*` helper hands back.
  3
  4A capture is a read-only sequence of what happened inside its block, in the
  5order it happened. It is complete when the block ends, so that is when it can
  6be read.
  7
  8`Captured` is what a test reads. `CaptureSource` is for writing a
  9`capture_*` helper whose captures all read from one place, such as the span
 10exporter or a connection's query log.
 11"""
 12
 13from collections.abc import Callable, Generator, Iterable, Iterator, Sequence
 14from contextlib import contextmanager
 15from typing import overload
 16
 17__all__ = ["CaptureSource", "Captured"]
 18
 19
 20class Captured[T](Sequence[T]):
 21    """
 22    What a `capture_*` block recorded, in the order it happened.
 23
 24        with capture_spans() as spans:
 25            Client().get("/")
 26
 27        assert len(spans) == 1
 28        assert spans[0].name == "GET /"
 29
 30    It's a sequence, so `len()`, indexing, slicing, iteration, `in` and
 31    truthiness all work, and nothing can be added to it or removed from it.
 32    It compares equal to a list or tuple holding the same items.
 33
 34    Read it after the block. Inside the block the capture is still going — a
 35    span that hasn't ended or a metric that hasn't been collected isn't there
 36    yet — so reading it raises instead of answering with part of the story.
 37    """
 38
 39    def __init__(self, *, helper: str) -> None:
 40        # The helper's name, for the messages: "capture_spans".
 41        self._helper = helper
 42        self._items: tuple[T, ...] | None = None
 43
 44    def finish(self, items: Iterable[T]) -> None:
 45        """
 46        Say what was captured. The helper that made this capture calls it
 47        when its block ends, and from then on the capture can be read.
 48        """
 49        self._items = tuple(items)
 50
 51    @property
 52    def finished(self) -> bool:
 53        """Whether the block has ended, so the capture can be read."""
 54        return self._items is not None
 55
 56    def _finished_items(self) -> tuple[T, ...]:
 57        if self._items is None:
 58            raise RuntimeError(
 59                f"{self._helper}() is still capturing — read what it captured"
 60                f" after the `with {self._helper}()` block ends, not inside it."
 61            )
 62        return self._items
 63
 64    def __len__(self) -> int:
 65        return len(self._finished_items())
 66
 67    @overload
 68    def __getitem__(self, index: int) -> T: ...
 69
 70    @overload
 71    def __getitem__(self, index: slice) -> Sequence[T]: ...
 72
 73    def __getitem__(self, index: int | slice) -> T | Sequence[T]:
 74        return self._finished_items()[index]
 75
 76    def __iter__(self) -> Iterator[T]:
 77        return iter(self._finished_items())
 78
 79    def __eq__(self, other: object) -> bool:
 80        # Equal to a list or tuple of the same items, so `captured == []`
 81        # says what it looks like it says instead of always being False.
 82        if isinstance(other, Captured):
 83            return self._finished_items() == other._finished_items()
 84        if isinstance(other, list | tuple):
 85            return list(self._finished_items()) == list(other)
 86        return NotImplemented
 87
 88    def __repr__(self) -> str:
 89        if self._items is None:
 90            return f"<{type(self).__name__}: still capturing>"
 91        return f"<{type(self).__name__} {list(self._items)!r}>"
 92
 93
 94class CaptureSource[T]:
 95    """
 96    The one place every capture of a kind reads from: the span exporter, a
 97    connection's query log. It is a list that grows as things happen, given
 98    here as a function that reads it and a function that empties it.
 99
100        _span_source = CaptureSource(
101            read=exporter.get_finished_spans, clear=exporter.clear
102        )
103
104        @contextmanager
105        def capture_spans():
106            captured = CapturedSpans()
107            with _span_source.capturing_into(captured):
108                yield captured
109
110    Captures nest: a project lifecycle can capture around every test, with
111    the test's own capture inside it. So a capture doesn't empty the list
112    when it starts. It remembers how long the list was, and gets what was
113    added after that. The list is emptied when no capture is open: before
114    the outermost one starts, and after it ends, so that nothing captured is
115    kept for the rest of the run.
116    """
117
118    def __init__(
119        self, *, read: Callable[[], Sequence[T]], clear: Callable[[], None]
120    ) -> None:
121        self._read = read
122        self._clear = clear
123        self._open_captures = 0
124
125    @property
126    def capturing(self) -> bool:
127        """Whether any capture that reads from here is open."""
128        return self._open_captures > 0
129
130    @contextmanager
131    def capturing_into(self, captured: Captured[T]) -> Generator[None]:
132        """
133        Capture for the duration of the block, and finish `captured` with
134        what was added during it. A block that raises still finishes its
135        capture.
136        """
137        if not self.capturing:
138            self._clear()
139        start = len(self._read())
140
141        self._open_captures += 1
142        try:
143            yield
144        finally:
145            self._open_captures -= 1
146            captured.finish(self._read()[start:])
147            if not self.capturing:
148                self._clear()