Add/fix types to a bunch of non-generated protobuf files (#7017)

stubs/protobuf/@tests/stubtest_allowlist.txt

@@ -2,34 +2,8 @@
 # carefully in mypy-protobuf which internally runs stubtest. Skip those here.
 google.protobuf\..*_pb2\..*
 
-google.protobuf.__version__
-google.protobuf.descriptor.Descriptor.__new__
-google.protobuf.descriptor.FieldDescriptor.__init__
-google.protobuf.descriptor.FieldDescriptor.__new__
-google.protobuf.descriptor.ServiceDescriptor.__new__
-google.protobuf.descriptor._NestedDescriptorBase.GetTopLevelContainingType
 google.protobuf.internal.containers.BaseContainer.sort
-google.protobuf.internal.containers.MessageMap.__init__
-google.protobuf.internal.containers.MessageMap.get
-google.protobuf.internal.containers.RepeatedCompositeFieldContainer.__init__
-google.protobuf.internal.containers.RepeatedScalarFieldContainer.__init__
-google.protobuf.internal.containers.ScalarMap.__init__
-google.protobuf.internal.containers.ScalarMap.get
-google.protobuf.internal.decoder.BytesDecoder
-google.protobuf.internal.decoder.EnumDecoder
-google.protobuf.internal.decoder.MessageSetItemDecoder
-google.protobuf.internal.decoder.StringDecoder
-google.protobuf.internal.encoder.MapSizer
-google.protobuf.internal.enum_type_wrapper.EnumTypeWrapper.DESCRIPTOR
-google.protobuf.internal.python_message.GeneratedProtocolMessageType.__init__
-google.protobuf.internal.python_message.GeneratedProtocolMessageType.__new__
-google.protobuf.internal.well_known_types.Any.Pack
-google.protobuf.internal.well_known_types.Error
-google.protobuf.internal.well_known_types.ParseError
-google.protobuf.message.Message.DESCRIPTOR
-google.protobuf.message.Message.Extensions
 google.protobuf.message.Message.SerializePartialToString
 google.protobuf.message.Message.SerializeToString
-google.protobuf.message.Message.__init__
 google.protobuf.service.Service.GetDescriptor
 google.protobuf.text_format.MessageToBytes

stubs/protobuf/google/protobuf/__init__.pyi

@@ -1 +1 @@
-__version__: bytes
+__version__: str

stubs/protobuf/google/protobuf/descriptor.pyi

@@ -42,20 +42,19 @@ class _NestedDescriptorBase(DescriptorBase):
         serialized_end=...,
         serialized_options=...,
     ) -> None: ...
-    def GetTopLevelContainingType(self): ...
     def CopyToProto(self, proto): ...
 
 class Descriptor(_NestedDescriptorBase):
     def __new__(
         cls,
-        name,
-        full_name,
-        filename,
-        containing_type,
-        fields,
-        nested_types,
-        enum_types,
-        extensions,
+        name=...,
+        full_name=...,
+        filename=...,
+        containing_type=...,
+        fields=...,
+        nested_types=...,
+        enum_types=...,
+        extensions=...,
         options=...,
         serialized_options=...,
         is_extendable=...,
@@ -162,9 +161,10 @@ class FieldDescriptor(DescriptorBase):
         extension_scope,
         options=...,
         serialized_options=...,
-        file=...,
         has_default_value=...,
         containing_oneof=...,
+        json_name=...,
+        file=...,
         create_key=...,
     ): ...
     name: Any
@@ -199,9 +199,10 @@ class FieldDescriptor(DescriptorBase):
         extension_scope,
         options=...,
         serialized_options=...,
-        file=...,
         has_default_value=...,
         containing_oneof=...,
+        json_name=...,
+        file=...,
         create_key=...,
     ) -> None: ...
     @staticmethod
@@ -267,10 +268,10 @@ class OneofDescriptor:
 class ServiceDescriptor(_NestedDescriptorBase):
     def __new__(
         cls,
-        name,
-        full_name,
-        index,
-        methods,
+        name=...,
+        full_name=...,
+        index=...,
+        methods=...,
         options=...,
         serialized_options=...,
         file=...,

stubs/protobuf/google/protobuf/internal/containers.pyi

@@ -5,6 +5,7 @@ from typing_extensions import SupportsIndex
 from google.protobuf.descriptor import Descriptor
 from google.protobuf.internal.message_listener import MessageListener
 from google.protobuf.internal.python_message import GeneratedProtocolMessageType
+from google.protobuf.internal.type_checkers import _ValueChecker
 from google.protobuf.message import Message
 
 _T = TypeVar("_T")
@@ -25,7 +26,7 @@ class BaseContainer(Sequence[_T]):
     def __getitem__(self, key: slice) -> list[_T]: ...
 
 class RepeatedScalarFieldContainer(BaseContainer[_ScalarV]):
-    def __init__(self, message_listener: MessageListener, message_descriptor: Descriptor) -> None: ...
+    def __init__(self, message_listener: MessageListener, type_checker: _ValueChecker[_ScalarV]) -> None: ...
     def append(self, value: _ScalarV) -> None: ...
     def insert(self, key: int, value: _ScalarV) -> None: ...
     def extend(self, elem_seq: Optional[Iterable[_ScalarV]]) -> None: ...
@@ -43,7 +44,7 @@ class RepeatedScalarFieldContainer(BaseContainer[_ScalarV]):
     def __eq__(self, other: object) -> bool: ...
 
 class RepeatedCompositeFieldContainer(BaseContainer[_MessageV]):
-    def __init__(self, message_listener: MessageListener, type_checker: Any) -> None: ...
+    def __init__(self, message_listener: MessageListener, message_descriptor: Descriptor) -> None: ...
     def add(self, **kwargs: Any) -> _MessageV: ...
     def append(self, value: _MessageV) -> None: ...
     def insert(self, key: int, value: _MessageV) -> None: ...
@@ -57,23 +58,45 @@ class RepeatedCompositeFieldContainer(BaseContainer[_MessageV]):
     def __eq__(self, other: object) -> bool: ...
 
 class ScalarMap(MutableMapping[_K, _ScalarV]):
+    def __init__(
+        self,
+        message_listener: MessageListener,
+        key_checker: _ValueChecker[_K],
+        value_checker: _ValueChecker[_ScalarV],
+        entry_descriptor: Descriptor,
+    ) -> None: ...
     def __setitem__(self, k: _K, v: _ScalarV) -> None: ...
     def __delitem__(self, v: _K) -> None: ...
     def __getitem__(self, k: _K) -> _ScalarV: ...
     def __len__(self) -> int: ...
     def __iter__(self) -> Iterator[_K]: ...
     def __eq__(self, other: object) -> bool: ...
+    @overload
+    def get(self, key: _K, default: None = ...) -> _ScalarV: ...
+    @overload
+    def get(self, key: _K, default: Union[_ScalarV, _T]) -> Union[_ScalarV, _T]: ...
     def MergeFrom(self: _M, other: _M): ...
     def InvalidateIterators(self) -> None: ...
     def GetEntryClass(self) -> GeneratedProtocolMessageType: ...
 
 class MessageMap(MutableMapping[_K, _MessageV]):
+    def __init__(
+        self,
+        message_listener: MessageListener,
+        message_descriptor: Descriptor,
+        key_checker: _ValueChecker[_K],
+        entry_descriptor: Descriptor,
+    ) -> None: ...
     def __setitem__(self, k: _K, v: _MessageV) -> None: ...
     def __delitem__(self, v: _K) -> None: ...
     def __getitem__(self, k: _K) -> _MessageV: ...
     def __len__(self) -> int: ...
     def __iter__(self) -> Iterator[_K]: ...
     def __eq__(self, other: object) -> bool: ...
+    @overload
+    def get(self, key: _K, default: None = ...) -> _MessageV: ...
+    @overload
+    def get(self, key: _K, default: Union[_MessageV, _T]) -> Union[_MessageV, _T]: ...
     def get_or_create(self, key: _K) -> _MessageV: ...
     def MergeFrom(self: _M, other: _M): ...
     def InvalidateIterators(self) -> None: ...

stubs/protobuf/google/protobuf/internal/decoder.pyi

@@ -1,30 +1,61 @@
-from typing import Any
+from typing import Any, Callable
+
+from google.protobuf.descriptor import Descriptor, FieldDescriptor
+from google.protobuf.message import Message
+
+_Decoder = Callable[[str, int, int, Message, dict[FieldDescriptor, Any]], int]
+_NewDefault = Callable[[Message], Message]
 
 def ReadTag(buffer, pos): ...
-def EnumDecoder(field_number, is_repeated, is_packed, key, new_default): ...
 
-Int32Decoder: Any
-Int64Decoder: Any
-UInt32Decoder: Any
-UInt64Decoder: Any
-SInt32Decoder: Any
-SInt64Decoder: Any
-Fixed32Decoder: Any
-Fixed64Decoder: Any
-SFixed32Decoder: Any
-SFixed64Decoder: Any
-FloatDecoder: Any
-DoubleDecoder: Any
-BoolDecoder: Any
+Int32Decoder: _Decoder
+Int64Decoder: _Decoder
+UInt32Decoder: _Decoder
+UInt64Decoder: _Decoder
+SInt32Decoder: _Decoder
+SInt64Decoder: _Decoder
+Fixed32Decoder: _Decoder
+Fixed64Decoder: _Decoder
+SFixed32Decoder: _Decoder
+SFixed64Decoder: _Decoder
+FloatDecoder: _Decoder
+DoubleDecoder: _Decoder
+BoolDecoder: _Decoder
 
-def StringDecoder(field_number, is_repeated, is_packed, key, new_default): ...
-def BytesDecoder(field_number, is_repeated, is_packed, key, new_default): ...
-def GroupDecoder(field_number, is_repeated, is_packed, key, new_default): ...
-def MessageDecoder(field_number, is_repeated, is_packed, key, new_default): ...
+def EnumDecoder(
+    field_number: int,
+    is_repeated: bool,
+    is_packed: bool,
+    key: FieldDescriptor,
+    new_default: _NewDefault,
+    clear_if_default: bool = ...,
+) -> _Decoder: ...
+def StringDecoder(
+    field_number: int,
+    is_repeated: bool,
+    is_packed: bool,
+    key: FieldDescriptor,
+    new_default: _NewDefault,
+    clear_if_default: bool = ...,
+) -> _Decoder: ...
+def BytesDecoder(
+    field_number: int,
+    is_repeated: bool,
+    is_packed: bool,
+    key: FieldDescriptor,
+    new_default: _NewDefault,
+    clear_if_default: bool = ...,
+) -> _Decoder: ...
+def GroupDecoder(
+    field_number: int, is_repeated: bool, is_packed: bool, key: FieldDescriptor, new_default: _NewDefault
+) -> _Decoder: ...
+def MessageDecoder(
+    field_number: int, is_repeated: bool, is_packed: bool, key: FieldDescriptor, new_default: _NewDefault
+) -> _Decoder: ...
 
-MESSAGE_SET_ITEM_TAG: Any
+MESSAGE_SET_ITEM_TAG: bytes
 
-def MessageSetItemDecoder(extensions_by_number): ...
-def MapDecoder(field_descriptor, new_default, is_message_map): ...
+def MessageSetItemDecoder(descriptor: Descriptor) -> _Decoder: ...
+def MapDecoder(field_descriptor, new_default, is_message_map) -> _Decoder: ...
 
 SkipField: Any

stubs/protobuf/google/protobuf/internal/encoder.pyi

@@ -1,34 +1,40 @@
-from typing import Any
+from typing import Callable
 
-Int32Sizer: Any
-UInt32Sizer: Any
-SInt32Sizer: Any
-Fixed32Sizer: Any
-Fixed64Sizer: Any
-BoolSizer: Any
+from google.protobuf.descriptor import FieldDescriptor
 
-def StringSizer(field_number, is_repeated, is_packed): ...
-def BytesSizer(field_number, is_repeated, is_packed): ...
-def GroupSizer(field_number, is_repeated, is_packed): ...
-def MessageSizer(field_number, is_repeated, is_packed): ...
-def MessageSetItemSizer(field_number): ...
-def MapSizer(field_descriptor): ...
-def TagBytes(field_number, wire_type): ...
+_Sizer = Callable[[int, bool, bool], int]
 
-Int32Encoder: Any
-UInt32Encoder: Any
-SInt32Encoder: Any
-Fixed32Encoder: Any
-Fixed64Encoder: Any
-SFixed32Encoder: Any
-SFixed64Encoder: Any
-FloatEncoder: Any
-DoubleEncoder: Any
+Int32Sizer: _Sizer
+UInt32Sizer: _Sizer
+SInt32Sizer: _Sizer
+Fixed32Sizer: _Sizer
+Fixed64Sizer: _Sizer
+BoolSizer: _Sizer
 
-def BoolEncoder(field_number, is_repeated, is_packed): ...
-def StringEncoder(field_number, is_repeated, is_packed): ...
-def BytesEncoder(field_number, is_repeated, is_packed): ...
-def GroupEncoder(field_number, is_repeated, is_packed): ...
-def MessageEncoder(field_number, is_repeated, is_packed): ...
-def MessageSetItemEncoder(field_number): ...
-def MapEncoder(field_descriptor): ...
+def StringSizer(field_number: int, is_repeated: bool, is_packed: bool) -> _Sizer: ...
+def BytesSizer(field_number: int, is_repeated: bool, is_packed: bool) -> _Sizer: ...
+def GroupSizer(field_number: int, is_repeated: bool, is_packed: bool) -> _Sizer: ...
+def MessageSizer(field_number: int, is_repeated: bool, is_packed: bool) -> _Sizer: ...
+def MessageSetItemSizer(field_number: int) -> _Sizer: ...
+def MapSizer(field_descriptor: FieldDescriptor, is_message_map: bool) -> _Sizer: ...
+def TagBytes(field_number: int, wire_type: int) -> bytes: ...
+
+_Encoder = Callable[[Callable[[bytes], int], bytes, bool], int]
+
+Int32Encoder: _Encoder
+UInt32Encoder: _Encoder
+SInt32Encoder: _Encoder
+Fixed32Encoder: _Encoder
+Fixed64Encoder: _Encoder
+SFixed32Encoder: _Encoder
+SFixed64Encoder: _Encoder
+FloatEncoder: _Encoder
+DoubleEncoder: _Encoder
+
+def BoolEncoder(field_number: int, is_repeated: bool, is_packed: bool) -> _Encoder: ...
+def StringEncoder(field_number: int, is_repeated: bool, is_packed: bool) -> _Encoder: ...
+def BytesEncoder(field_number: int, is_repeated: bool, is_packed: bool) -> _Encoder: ...
+def GroupEncoder(field_number: int, is_repeated: bool, is_packed: bool) -> _Encoder: ...
+def MessageEncoder(field_number: int, is_repeated: bool, is_packed: bool) -> _Encoder: ...
+def MessageSetItemEncoder(field_number: int) -> _Encoder: ...
+def MapEncoder(field_descriptor: FieldDescriptor) -> _Encoder: ...

stubs/protobuf/google/protobuf/internal/enum_type_wrapper.pyi

@@ -7,7 +7,7 @@ _V = TypeVar("_V", bound=int)
 # Expose a generic version so that those using mypy-protobuf
 # can get autogenerated NewType wrapper around the int values
 class _EnumTypeWrapper(Generic[_V]):
-    DESCRIPTOR: EnumDescriptor
+    DESCRIPTOR: EnumDescriptor | None
     def __init__(self, enum_type: EnumDescriptor) -> None: ...
     def Name(self, number: _V) -> str: ...
     def Value(self, name: Text | bytes) -> _V: ...

stubs/protobuf/google/protobuf/internal/python_message.pyi

@@ -1 +1,3 @@
-class GeneratedProtocolMessageType(type): ...
+class GeneratedProtocolMessageType(type):
+    def __new__(cls, name, bases, dictionary): ...
+    def __init__(cls, name, bases, dictionary): ...

stubs/protobuf/google/protobuf/internal/type_checkers.pyi

@@ -0,0 +1,15 @@
+from typing import Generic, Protocol, TypeVar
+
+_T = TypeVar("_T")
+
+class _ValueChecker(Protocol[_T]):
+    def CheckValue(self, proposed_value: _T) -> _T: ...
+    def DefaultValue(self) -> _T: ...
+
+class TypeChecker(Generic[_T]):
+    def __init__(self, *acceptable_types: _T): ...
+    def CheckValue(self, proposed_value: _T) -> _T: ...
+
+class TypeCheckerWithDefault(TypeChecker[_T]):
+    def __init__(self, default_value: _T, *acceptable_types: _T): ...
+    def DefaultValue(self) -> _T: ...

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

@@ -1,13 +1,10 @@
 from datetime import datetime, timedelta
 from typing import Any as tAny, Optional
 
-class Error(Exception): ...
-class ParseError(Error): ...
-
 class Any:
     type_url: tAny = ...
     value: tAny = ...
-    def Pack(self, msg: tAny, type_url_prefix: bytes = ..., deterministic: Optional[tAny] = ...) -> None: ...
+    def Pack(self, msg: tAny, type_url_prefix: str = ..., deterministic: Optional[tAny] = ...) -> None: ...
     def Unpack(self, msg: tAny): ...
     def TypeName(self): ...
     def Is(self, descriptor: tAny): ...

stubs/protobuf/google/protobuf/message.pyi

@@ -2,7 +2,7 @@ from _typeshed import Self
 from typing import Any, Sequence, TypeVar
 
 from .descriptor import Descriptor, FieldDescriptor
-from .internal.extension_dict import _ExtensionDict, _ExtensionFieldDescriptor
+from .internal.extension_dict import _ExtensionFieldDescriptor
 
 class Error(Exception): ...
 class DecodeError(Error): ...
@@ -11,7 +11,7 @@ class EncodeError(Error): ...
 _M = TypeVar("_M", bound=Message)  # message type (of self)
 
 class Message:
-    DESCRIPTOR: Descriptor
+    DESCRIPTOR: Descriptor | None
     def __deepcopy__(self, memo=...): ...
     def __eq__(self, other_msg): ...
     def __ne__(self, other_msg): ...
@@ -31,9 +31,6 @@ class Message:
     def ByteSize(self) -> int: ...
     @classmethod
     def FromString(cls: type[Self], s: bytes) -> Self: ...
-    # The TypeVar must be bound to `Message` or we get mypy errors, so we cannot use `Self` for `Extensions`
-    @property
-    def Extensions(self: _M) -> _ExtensionDict[_M]: ...
     # Intentionally left out typing on these three methods, because they are
     # stringly typed and it is not useful to call them on a Message directly.
     # We prefer more specific typing on individual subclasses of Message
@@ -42,4 +39,4 @@ class Message:
     def ClearField(self, field_name: Any) -> None: ...
     def WhichOneof(self, oneof_group: Any) -> Any: ...
     # TODO: check kwargs
-    def __init__(self, **kwargs) -> None: ...
+    def __init__(self, *args, **kwargs) -> None: ...

tests/pytype_exclude_list.txt

@@ -30,6 +30,8 @@ stubs/protobuf/google/protobuf/duration_pb2.pyi
 stubs/protobuf/google/protobuf/empty_pb2.pyi
 stubs/protobuf/google/protobuf/field_mask_pb2.pyi
 stubs/protobuf/google/protobuf/internal/containers.pyi
+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/json_format.pyi