protobuf: Annotate well_known_types.pyi (#9323)

Co-authored-by: Alex Waygood <[email protected]>

Author: ikonst

stubs/protobuf/@tests/stubtest_allowlist.txt

@@ -36,3 +36,6 @@ google.protobuf.internal.containers.BaseContainer.__hash__
 
 # Metaclass differs:
 google.protobuf.descriptor.OneofDescriptor
+
+# Runtime does not have __iter__ (yet...): hack in spirit of https://github.com/python/typeshed/issues/7813
+google.protobuf.internal.well_known_types.ListValue.__iter__

stubs/protobuf/@tests/test_cases/check_struct.py

@@ -0,0 +1,18 @@
+# pyright: reportUnnecessaryTypeIgnoreComment=true
+from __future__ import annotations
+
+from google.protobuf.struct_pb2 import ListValue, Struct
+
+list_value = ListValue()
+
+lst = list(list_value)  # Ensure type checkers recognise that the class is iterable (doesn't have an `__iter__` method at runtime)
+
+list_value[0] = 42.42
+list_value[0] = "42"
+list_value[0] = None
+list_value[0] = True
+list_value[0] = [42.42, "42", None, True, [42.42, "42", None, True], {"42": 42}]
+list_value[0] = ListValue()
+list_value[0] = Struct()
+
+list_element = list_value[0]

stubs/protobuf/google/protobuf/internal/well_known_types.pyi

@@ -1,5 +1,10 @@
+from _typeshed import SupportsItems
+from collections.abc import Iterable, Iterator, KeysView, Mapping, Sequence
 from datetime import datetime, timedelta, tzinfo
 from typing import Any as tAny
+from typing_extensions import TypeAlias
+
+from google.protobuf import struct_pb2
 
 class Any:
     type_url: tAny = ...
@@ -63,29 +68,34 @@ class _FieldMaskTree:
     def AddLeafNodes(self, prefix: tAny, node: tAny) -> None: ...
     def MergeMessage(self, source: tAny, destination: tAny, replace_message: tAny, replace_repeated: tAny) -> None: ...
 
+_StructValue: TypeAlias = struct_pb2.Struct | struct_pb2.ListValue | str | float | bool | None
+_StructValueArg: TypeAlias = _StructValue | Mapping[str, _StructValueArg] | Sequence[_StructValueArg]
+
 class Struct:
-    def __getitem__(self, key: tAny): ...
-    def __contains__(self, item: tAny): ...
-    def __setitem__(self, key: tAny, value: tAny) -> None: ...
-    def __delitem__(self, key: tAny) -> None: ...
+    def __getitem__(self, key: str) -> _StructValue: ...
+    def __contains__(self, item: object) -> bool: ...
+    def __setitem__(self, key: str, value: _StructValueArg) -> None: ...
+    def __delitem__(self, key: str) -> None: ...
     def __len__(self) -> int: ...
-    def __iter__(self): ...
-    def keys(self): ...
-    def values(self): ...
-    def items(self): ...
-    def get_or_create_list(self, key: tAny): ...
-    def get_or_create_struct(self, key: tAny): ...
-    def update(self, dictionary: tAny) -> None: ...
+    def __iter__(self) -> Iterator[str]: ...
+    def keys(self) -> KeysView[str]: ...
+    def values(self) -> list[_StructValue]: ...
+    def items(self) -> list[tuple[str, _StructValue]]: ...
+    def get_or_create_list(self, key: str) -> struct_pb2.ListValue: ...
+    def get_or_create_struct(self, key: str) -> struct_pb2.Struct: ...
+    def update(self, dictionary: SupportsItems[str, _StructValueArg]) -> None: ...
 
 class ListValue:
     def __len__(self) -> int: ...
-    def append(self, value: tAny) -> None: ...
-    def extend(self, elem_seq: tAny) -> None: ...
-    def __getitem__(self, index: tAny): ...
-    def __setitem__(self, index: tAny, value: tAny) -> None: ...
-    def __delitem__(self, key: tAny) -> None: ...
-    def items(self) -> None: ...
-    def add_struct(self): ...
-    def add_list(self): ...
+    def append(self, value: _StructValue) -> None: ...
+    def extend(self, elem_seq: Iterable[_StructValue]) -> None: ...
+    def __getitem__(self, index: int) -> _StructValue: ...
+    def __setitem__(self, index: int, value: _StructValueArg) -> None: ...
+    def __delitem__(self, key: int) -> None: ...
+    # Doesn't actually exist at runtime; needed so type checkers understand the class is iterable
+    def __iter__(self) -> Iterator[_StructValue]: ...
+    def items(self) -> Iterator[_StructValue]: ...
+    def add_struct(self) -> struct_pb2.Struct: ...
+    def add_list(self) -> struct_pb2.ListValue: ...
 
 WKTBASES: dict[str, type[tAny]]

tests/pytype_exclude_list.txt

@@ -27,6 +27,7 @@ stubs/protobuf/google/protobuf/internal/decoder.pyi
 stubs/protobuf/google/protobuf/internal/encoder.pyi
 stubs/protobuf/google/protobuf/internal/enum_type_wrapper.pyi
 stubs/protobuf/google/protobuf/internal/extension_dict.pyi
+stubs/protobuf/google/protobuf/internal/well_known_types.pyi
 stubs/protobuf/google/protobuf/json_format.pyi
 stubs/protobuf/google/protobuf/message.pyi
 stubs/protobuf/google/protobuf/message_factory.pyi