Module astrapy.event_observers.context_managers

@contextmanager
def event_collector(
    target: DB_OBJ,
    *,
    destination: list[ObservableEvent]
    | dict[ObservableEventType, list[ObservableEvent]],
    event_types: Iterable[ObservableEventType] | None = None,
) -> Iterator[DB_OBJ]:
    """
    Create a context manager wrapping a list (of `ObservableEvent` objects) or
    a dict (with `ObservableEventType` keys and `list[ObservableEvent]` values),
    for quickly instrumenting a client, database, table, collection or admin object.

    This utility is meant to be used in a `with` statement, so that events emitted
    by the instrumented classes within the with block are captured into the provided
    lists/dictionaries. Events emitted by spawned classes (e.g. when a Database
    creates a Collection in the `with` block) are also received this way.

    Once outside the `with` block, collection of events stops, but the provided
    destination can still be accessed according to its ordinary scoping.

    Args:
        target: an object that issues Data API / DevOps API requests. The target
            must have a `with_options` suitable method: meaning, it can be any of
            the following: `DataAPIClient`, `AsyncDatabase`, `Database`,
            `AsyncCollection`, `Collection`, `AsyncTable`, `Table`, `AstraDBAdmin`,
            `AstraDBDatabaseAdmin`, `DataAPIDatabaseAdmin`.
        destination: a list or a dictionary where the collected events will be stored.
            For dictionaries, events are grouped into lists, one per each event type,
            stored under the corresponding `ObservableEventType` value as dict key.
        event_types: if provided, it's a list of event types so that only
            events matching this filter are processed.

    Returns:
        Yields an instrumented version of the input target, with an added observer
        set to accumulate the received events into the provided destination. Any
        pre-existing observer is untouched.

    Example:
        >>> ev_lst: list[ObservableEvent] = []
        >>> with event_collector(db, destination=ev_lst) as instrumented_db:
        ...     _ = instrumented_db.list_table_names()
        ...     table = instrumented_db.get_table("my_table")
        ...     _ = table.find_one({"k": 101})
        ...
        >>> print(len(ev_lst))
        5
        >>> print(ev_lst[0].event_type)
        ObservableEventType.REQUEST
    """
    observer_id_ = f"observer_{str(uuid7())}"
    observer_: Observer
    if isinstance(destination, list):
        observer_ = Observer.from_event_list(destination, event_types=event_types)
    else:
        observer_ = Observer.from_event_dict(destination, event_types=event_types)
    api_options = APIOptions(event_observers={observer_id_: observer_})
    target_ = target.with_options(api_options=api_options)
    try:
        yield target_
    finally:
        observer_.enabled = False

class OptionAwareDatabaseObject(Protocol):
    def with_options(
        self,
        *,
        api_options: APIOptions | UnsetType = _UNSET,
    ) -> Self: ...