xref: /aosp_15_r20/external/grpc-grpc/src/python/grpcio/grpc/aio/_metadata.py (revision cc02d7e222339f7a4f6ba5f422e6413f4bd931f2)
1# Copyright 2020 gRPC authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14"""Implementation of the metadata abstraction for gRPC Asyncio Python."""
15from collections import OrderedDict
16from collections import abc
17from typing import Any, Iterator, List, Optional, Tuple, Union
18
19MetadataKey = str
20MetadataValue = Union[str, bytes]
21
22
23class Metadata(abc.Collection):
24    """Metadata abstraction for the asynchronous calls and interceptors.
25
26    The metadata is a mapping from str -> List[str]
27
28    Traits
29        * Multiple entries are allowed for the same key
30        * The order of the values by key is preserved
31        * Getting by an element by key, retrieves the first mapped value
32        * Supports an immutable view of the data
33        * Allows partial mutation on the data without recreating the new object from scratch.
34    """
35
36    def __init__(self, *args: Tuple[MetadataKey, MetadataValue]) -> None:
37        self._metadata = OrderedDict()
38        for md_key, md_value in args:
39            self.add(md_key, md_value)
40
41    @classmethod
42    def from_tuple(cls, raw_metadata: tuple):
43        if raw_metadata:
44            return cls(*raw_metadata)
45        return cls()
46
47    def add(self, key: MetadataKey, value: MetadataValue) -> None:
48        self._metadata.setdefault(key, [])
49        self._metadata[key].append(value)
50
51    def __len__(self) -> int:
52        """Return the total number of elements that there are in the metadata,
53        including multiple values for the same key.
54        """
55        return sum(map(len, self._metadata.values()))
56
57    def __getitem__(self, key: MetadataKey) -> MetadataValue:
58        """When calling <metadata>[<key>], the first element of all those
59        mapped for <key> is returned.
60        """
61        try:
62            return self._metadata[key][0]
63        except (ValueError, IndexError) as e:
64            raise KeyError("{0!r}".format(key)) from e
65
66    def __setitem__(self, key: MetadataKey, value: MetadataValue) -> None:
67        """Calling metadata[<key>] = <value>
68        Maps <value> to the first instance of <key>.
69        """
70        if key not in self:
71            self._metadata[key] = [value]
72        else:
73            current_values = self.get_all(key)
74            self._metadata[key] = [value, *current_values[1:]]
75
76    def __delitem__(self, key: MetadataKey) -> None:
77        """``del metadata[<key>]`` deletes the first mapping for <key>."""
78        current_values = self.get_all(key)
79        if not current_values:
80            raise KeyError(repr(key))
81        self._metadata[key] = current_values[1:]
82
83    def delete_all(self, key: MetadataKey) -> None:
84        """Delete all mappings for <key>."""
85        del self._metadata[key]
86
87    def __iter__(self) -> Iterator[Tuple[MetadataKey, MetadataValue]]:
88        for key, values in self._metadata.items():
89            for value in values:
90                yield (key, value)
91
92    def keys(self) -> abc.KeysView:
93        return abc.KeysView(self)
94
95    def values(self) -> abc.ValuesView:
96        return abc.ValuesView(self)
97
98    def items(self) -> abc.ItemsView:
99        return abc.ItemsView(self)
100
101    def get(
102        self, key: MetadataKey, default: MetadataValue = None
103    ) -> Optional[MetadataValue]:
104        try:
105            return self[key]
106        except KeyError:
107            return default
108
109    def get_all(self, key: MetadataKey) -> List[MetadataValue]:
110        """For compatibility with other Metadata abstraction objects (like in Java),
111        this would return all items under the desired <key>.
112        """
113        return self._metadata.get(key, [])
114
115    def set_all(self, key: MetadataKey, values: List[MetadataValue]) -> None:
116        self._metadata[key] = values
117
118    def __contains__(self, key: MetadataKey) -> bool:
119        return key in self._metadata
120
121    def __eq__(self, other: Any) -> bool:
122        if isinstance(other, self.__class__):
123            return self._metadata == other._metadata
124        if isinstance(other, tuple):
125            return tuple(self) == other
126        return NotImplemented  # pytype: disable=bad-return-type
127
128    def __add__(self, other: Any) -> "Metadata":
129        if isinstance(other, self.__class__):
130            return Metadata(*(tuple(self) + tuple(other)))
131        if isinstance(other, tuple):
132            return Metadata(*(tuple(self) + other))
133        return NotImplemented  # pytype: disable=bad-return-type
134
135    def __repr__(self) -> str:
136        view = tuple(self)
137        return "{0}({1!r})".format(self.__class__.__name__, view)
138