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

nipunn1313 · nipunn1313 · commit 02a81c7dd020 · 2022-01-24T09:09:53.000-08:00
diff --git a/stubs/protobuf/@tests/stubtest_allowlist.txt b/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
diff --git a/stubs/protobuf/google/protobuf/__init__.pyi b/stubs/protobuf/google/protobuf/__init__.pyi
@@ -1 +1 @@
-__version__: bytes
+__version__: str
diff --git a/stubs/protobuf/google/protobuf/descriptor.pyi b/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=...,
diff --git a/stubs/protobuf/google/protobuf/internal/containers.pyi b/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 TypeChecker
 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: TypeChecker) -> 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: TypeChecker,
+        value_checker: TypeChecker,
+        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: TypeChecker,
+        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: ...
diff --git a/stubs/protobuf/google/protobuf/internal/decoder.pyi b/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
diff --git a/stubs/protobuf/google/protobuf/internal/encoder.pyi b/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: ...
diff --git a/stubs/protobuf/google/protobuf/internal/enum_type_wrapper.pyi b/stubs/protobuf/google/protobuf/internal/enum_type_wrapper.pyi
@@ -1,4 +1,4 @@
-from typing import Generic, Text, TypeVar
+from typing import Generic, Optional, Text, TypeVar
 
 from google.protobuf.descriptor import EnumDescriptor
 
@@ -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: Optional[EnumDescriptor]
     def __init__(self, enum_type: EnumDescriptor) -> None: ...
     def Name(self, number: _V) -> str: ...
     def Value(self, name: Text | bytes) -> _V: ...
diff --git a/stubs/protobuf/google/protobuf/internal/python_message.pyi b/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): ...
diff --git a/stubs/protobuf/google/protobuf/internal/type_checkers.pyi b/stubs/protobuf/google/protobuf/internal/type_checkers.pyi
@@ -0,0 +1,7 @@
+from typing import Any, TypeVar
+
+_T = TypeVar("_T")
+
+class TypeChecker:
+    def __init__(self, *acceptable_types: Any): ...
+    def CheckValue(self, proposed_value: _T) -> _T: ...
diff --git a/stubs/protobuf/google/protobuf/internal/well_known_types.pyi b/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): ...
diff --git a/stubs/protobuf/google/protobuf/message.pyi b/stubs/protobuf/google/protobuf/message.pyi
diff --git a/tests/pytype_exclude_list.txt b/tests/pytype_exclude_list.txt
