hermes-chatto-plugin/vendor/connectrpc/_headers.py

from __future__ import annotations

from collections.abc import (
    ItemsView,
    Iterator,
    KeysView,
    Mapping,
    MutableMapping,
    Sequence,
    ValuesView,
)
from typing import cast


class Headers(MutableMapping[str, str]):
    """Container of HTTP headers.

    This class behaves like a dictionary with case-insensitive keys and
    string values. It additionally can be used to store multiple values for
    the same key by using the `add` method, and retrieve a view including
    duplicates using `allitems`.
    """

    _store: dict[str, str]
    _extra: dict[str, list[str]] | None

    def __init__(
        self, items: Mapping[str, str] | Sequence[tuple[str, str]] = ()
    ) -> None:
        self._extra = None
        if isinstance(items, Mapping):
            # ty does not preserve type parameters when narrowing a union via
            # `isinstance(x, Mapping)`, so we re-assert `Mapping[str, str]`.
            # See https://github.com/astral-sh/ty/issues/456.
            items = cast("Mapping[str, str]", items)
            self._store = {k.lower(): v for k, v in items.items()}
        else:
            self._store = {}
            for k, v in items:
                self.add(k, v)

    def __getitem__(self, key: str) -> str:
        key = key.lower()
        return self._store[key]

    def __setitem__(self, key: str, value: str) -> None:
        key = key.lower()
        self._store[key] = value
        if self._extra is not None:
            self._extra.pop(key, None)

    def __delitem__(self, key: str) -> None:
        key = key.lower()
        del self._store[key]
        # If value wasn't in store, it's not in extra
        # so no need to worry about exception.
        if self._extra is not None:
            self._extra.pop(key, None)

    def __iter__(self) -> Iterator[str]:
        return iter(self._store)

    def __len__(self) -> int:
        return len(self._store)

    def __repr__(self) -> str:
        return repr(list(self.allitems()))

    def add(self, key: str, value: str) -> None:
        """Add a header, appending to existing values without overwriting.

        To overwrite an existing value, use `self[key] = value` instead.
        """
        key = key.lower()
        if key in self._store:
            if self._extra is None:
                self._extra = {}
            self._extra.setdefault(key, []).append(value)
            return
        self._store[key] = value

    def clear(self) -> None:
        """Clear all headers."""
        self._store.clear()
        self._extra = None

    def getall(self, key: str) -> Sequence[str]:
        """Get all values for a header key, including duplicates."""
        key = key.lower()
        if key not in self._store:
            return ()
        res = [self._store[key]]
        if self._extra is not None:
            res.extend(self._extra.get(key, ()))
        return res

    def allitems(self) -> ItemsView[str, str]:
        """Return an iterable view of all header items, including duplicates."""
        if self._extra is None:
            return self._store.items()
        return _AllItemsView(self._store, self._extra)

    # Commonly used functions that delegate to _store for performance vs the base class
    # implementations that rely on dunder methods.

    def items(self) -> ItemsView[str, str]:
        """Return an iterable view of the headers, without duplicates."""
        return self._store.items()

    def keys(self) -> KeysView[str]:
        """Return an iterable view of the header keys."""
        return self._store.keys()

    def values(self) -> ValuesView[str]:
        """Return an iterable view of the header values, without duplicates."""
        return self._store.values()

    def __contains__(self, key: object) -> bool:
        if not isinstance(key, str):
            return False
        key = key.lower()
        return key in self._store


class _AllItemsView(ItemsView[str, str]):
    """An iterable view of all header items, including duplicates."""

    _store: dict[str, str]
    _extra: dict[str, list[str]] | None

    def __init__(
        self, store: dict[str, str], extra: dict[str, list[str]] | None
    ) -> None:
        self._store = store
        self._extra = extra

    def __iter__(self) -> Iterator[tuple[str, str]]:
        for key, v in self._store.items():
            yield (key, v)
            if self._extra:
                for vv in self._extra.get(key, ()):
                    yield (key, vv)

    def __len__(self) -> int:
        size = len(self._store)
        if self._extra:
            size += sum(len(v) for v in self._extra.values())
        return size