Add/fix types to a bunch of non-generated protobuf files

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
-
-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): ...
-
-MESSAGE_SET_ITEM_TAG: Any
-
-def MessageSetItemDecoder(extensions_by_number): ...
-def MapDecoder(field_descriptor, new_default, is_message_map): ...
+
+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 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: bytes
+
+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): ...