Mark mangled globals as TypeAlias

nipunn1313 · Apr 25, 2024 · 0a10dbd · 0a10dbd
1 parent e44c1bc
commit 0a10dbd
Show file tree

Hide file tree

Showing 19 changed files with 109 additions and 34 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,5 +1,7 @@
 ## Upcoming
 
+- Mark top-level mangled identifiers `global___*` as `TypeAlias`.
+
 ## 3.6.0
 
 - Remove 3.7 compatibility for typing_extensions.final/Literal

diff --git a/README.md b/README.md
@@ -302,6 +302,7 @@ black .
 - [@Alphadelta14](https://github.com/Alphadelta14)
 - [@fergyfresh](https://github.com/fergyfresh)
 - [@AlexWaygood](https://github.com/AlexWaygood)
+- [@Avasam](https://github.com/Avasam)
 
 ## Licence etc.
 

diff --git a/mypy_protobuf/extensions_pb2.pyi b/mypy_protobuf/extensions_pb2.pyi
@@ -8,8 +8,14 @@ import google.protobuf.descriptor
 import google.protobuf.descriptor_pb2
 import google.protobuf.internal.extension_dict
 import google.protobuf.message
+import sys
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -34,7 +40,7 @@ class FieldOptions(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]) -> None: ...
 
-global___FieldOptions = FieldOptions
+global___FieldOptions: typing_extensions.TypeAlias = FieldOptions
 
 OPTIONS_FIELD_NUMBER: builtins.int
 CASTTYPE_FIELD_NUMBER: builtins.int

diff --git a/mypy_protobuf/main.py b/mypy_protobuf/main.py
@@ -376,7 +376,7 @@ def write_enums(
                 scl + [d.EnumDescriptorProto.VALUE_FIELD_NUMBER],
             )
             if prefix == "" and not self.readable_stubs:
-                wl(f"{_mangle_global_identifier(class_name)} = {class_name}")
+                wl(f"{_mangle_global_identifier(class_name)}: {self._import('typing_extensions', 'TypeAlias')} = {class_name}")
             wl("")
 
     def write_messages(
@@ -483,7 +483,7 @@ def write_messages(
 
             if prefix == "" and not self.readable_stubs:
                 wl("")
-                wl(f"{_mangle_global_identifier(class_name)} = {class_name}")
+                wl(f"{_mangle_global_identifier(class_name)}: {self._import('typing_extensions', 'TypeAlias')} = {class_name}")
             wl("")
 
     def write_stringly_typed_fields(self, desc: d.DescriptorProto) -> None:

diff --git a/test/generated/google/protobuf/duration_pb2.pyi b/test/generated/google/protobuf/duration_pb2.pyi
@@ -36,8 +36,14 @@ import builtins
 import google.protobuf.descriptor
 import google.protobuf.internal.well_known_types
 import google.protobuf.message
+import sys
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -127,4 +133,4 @@ class Duration(google.protobuf.message.Message, google.protobuf.internal.well_kn
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["nanos", b"nanos", "seconds", b"seconds"]) -> None: ...
 
-global___Duration = Duration
+global___Duration: typing_extensions.TypeAlias = Duration
diff --git a/test/generated/mypy_protobuf/extensions_pb2.pyi b/test/generated/mypy_protobuf/extensions_pb2.pyi
@@ -8,8 +8,14 @@ import google.protobuf.descriptor
 import google.protobuf.descriptor_pb2
 import google.protobuf.internal.extension_dict
 import google.protobuf.message
+import sys
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -34,7 +40,7 @@ class FieldOptions(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["casttype", b"casttype", "keytype", b"keytype", "valuetype", b"valuetype"]) -> None: ...
 
-global___FieldOptions = FieldOptions
+global___FieldOptions: typing_extensions.TypeAlias = FieldOptions
 
 OPTIONS_FIELD_NUMBER: builtins.int
 CASTTYPE_FIELD_NUMBER: builtins.int

diff --git a/test/generated/testproto/Capitalized/Capitalized_pb2.pyi b/test/generated/testproto/Capitalized/Capitalized_pb2.pyi
@@ -6,8 +6,14 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
+import sys
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -23,7 +29,7 @@ class lower(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["a", b"a"]) -> None: ...
 
-global___lower = lower
+global___lower: typing_extensions.TypeAlias = lower
 
 @typing.final
 class Upper(google.protobuf.message.Message):
@@ -40,7 +46,7 @@ class Upper(google.protobuf.message.Message):
     def HasField(self, field_name: typing.Literal["Lower", b"Lower"]) -> builtins.bool: ...
     def ClearField(self, field_name: typing.Literal["Lower", b"Lower"]) -> None: ...
 
-global___Upper = Upper
+global___Upper: typing_extensions.TypeAlias = Upper
 
 @typing.final
 class lower2(google.protobuf.message.Message):
@@ -57,4 +63,4 @@ class lower2(google.protobuf.message.Message):
     def HasField(self, field_name: typing.Literal["upper", b"upper"]) -> builtins.bool: ...
     def ClearField(self, field_name: typing.Literal["upper", b"upper"]) -> None: ...
 
-global___lower2 = lower2
+global___lower2: typing_extensions.TypeAlias = lower2
diff --git a/test/generated/testproto/comment_special_chars_pb2.pyi b/test/generated/testproto/comment_special_chars_pb2.pyi
@@ -6,8 +6,14 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
+import sys
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -73,4 +79,4 @@ class Test(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["a", b"a", "b", b"b", "c", b"c", "d", b"d", "e", b"e", "f", b"f", "g", b"g", "h", b"h", "i", b"i", "j", b"j", "k", b"k"]) -> None: ...
 
-global___Test = Test
+global___Test: typing_extensions.TypeAlias = Test
diff --git a/test/generated/testproto/dot/com/test_pb2.pyi b/test/generated/testproto/dot/com/test_pb2.pyi
@@ -6,8 +6,14 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
+import sys
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -23,4 +29,4 @@ class TestMessage(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["foo", b"foo"]) -> None: ...
 
-global___TestMessage = TestMessage
+global___TestMessage: typing_extensions.TypeAlias = TestMessage
diff --git a/test/generated/testproto/grpc/dummy_pb2.pyi b/test/generated/testproto/grpc/dummy_pb2.pyi
@@ -6,8 +6,14 @@ https://github.com/vmagamedov/grpclib/blob/master/tests/dummy.proto"""
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
+import sys
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -23,7 +29,7 @@ class DummyRequest(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["value", b"value"]) -> None: ...
 
-global___DummyRequest = DummyRequest
+global___DummyRequest: typing_extensions.TypeAlias = DummyRequest
 
 @typing.final
 class DummyReply(google.protobuf.message.Message):
@@ -38,4 +44,4 @@ class DummyReply(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["value", b"value"]) -> None: ...
 
-global___DummyReply = DummyReply
+global___DummyReply: typing_extensions.TypeAlias = DummyReply
diff --git a/test/generated/testproto/inner/inner_pb2.pyi b/test/generated/testproto/inner/inner_pb2.pyi
@@ -6,9 +6,15 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
+import sys
 import testproto.test3_pb2
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -24,4 +30,4 @@ class Inner(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["a", b"a"]) -> None: ...
 
-global___Inner = Inner
+global___Inner: typing_extensions.TypeAlias = Inner
diff --git a/test/generated/testproto/nested/nested_pb2.pyi b/test/generated/testproto/nested/nested_pb2.pyi
@@ -31,7 +31,7 @@ class Nested(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["a", b"a"]) -> None: ...
 
-global___Nested = Nested
+global___Nested: typing_extensions.TypeAlias = Nested
 
 @typing.final
 class AnotherNested(google.protobuf.message.Message):
@@ -93,4 +93,4 @@ class AnotherNested(google.protobuf.message.Message):
         self,
     ) -> None: ...
 
-global___AnotherNested = AnotherNested
+global___AnotherNested: typing_extensions.TypeAlias = AnotherNested
diff --git a/test/generated/testproto/nopackage_pb2.pyi b/test/generated/testproto/nopackage_pb2.pyi
@@ -8,8 +8,14 @@ import collections.abc
 import google.protobuf.descriptor
 import google.protobuf.internal.containers
 import google.protobuf.message
+import sys
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -22,7 +28,7 @@ class NoPackage(google.protobuf.message.Message):
         self,
     ) -> None: ...
 
-global___NoPackage = NoPackage
+global___NoPackage: typing_extensions.TypeAlias = NoPackage
 
 @typing.final
 class NoPackage2(google.protobuf.message.Message):
@@ -43,4 +49,4 @@ class NoPackage2(google.protobuf.message.Message):
     def HasField(self, field_name: typing.Literal["np", b"np"]) -> builtins.bool: ...
     def ClearField(self, field_name: typing.Literal["np", b"np", "np_rep", b"np_rep"]) -> None: ...
 
-global___NoPackage2 = NoPackage2
+global___NoPackage2: typing_extensions.TypeAlias = NoPackage2
diff --git a/test/generated/testproto/readme_enum_pb2.pyi b/test/generated/testproto/readme_enum_pb2.pyi
@@ -29,4 +29,4 @@ class MyEnum(_MyEnum, metaclass=_MyEnumEnumTypeWrapper): ...
 
 HELLO: MyEnum.ValueType  # 0
 WORLD: MyEnum.ValueType  # 1
-global___MyEnum = MyEnum
+global___MyEnum: typing_extensions.TypeAlias = MyEnum
diff --git a/test/generated/testproto/test3_pb2.pyi b/test/generated/testproto/test3_pb2.pyi
@@ -35,7 +35,7 @@ class OuterEnum(_OuterEnum, metaclass=_OuterEnumEnumTypeWrapper): ...
 UNKNOWN: OuterEnum.ValueType  # 0
 FOO3: OuterEnum.ValueType  # 1
 BAR3: OuterEnum.ValueType  # 2
-global___OuterEnum = OuterEnum
+global___OuterEnum: typing_extensions.TypeAlias = OuterEnum
 
 @typing.final
 class OuterMessage3(google.protobuf.message.Message):
@@ -50,7 +50,7 @@ class OuterMessage3(google.protobuf.message.Message):
     ) -> None: ...
     def ClearField(self, field_name: typing.Literal["a_string", b"a_string"]) -> None: ...
 
-global___OuterMessage3 = OuterMessage3
+global___OuterMessage3: typing_extensions.TypeAlias = OuterMessage3
 
 @typing.final
 class SimpleProto3(google.protobuf.message.Message):
@@ -206,4 +206,4 @@ class SimpleProto3(google.protobuf.message.Message):
     @typing.overload
     def WhichOneof(self, oneof_group: typing.Literal["b_oneof", b"b_oneof"]) -> typing.Literal["b_oneof_1", "b_oneof_2"] | None: ...
 
-global___SimpleProto3 = SimpleProto3
+global___SimpleProto3: typing_extensions.TypeAlias = SimpleProto3
diff --git a/test/generated/testproto/test_extensions2_pb2.pyi b/test/generated/testproto/test_extensions2_pb2.pyi
@@ -7,9 +7,15 @@ import builtins
 import google.protobuf.descriptor
 import google.protobuf.internal.extension_dict
 import google.protobuf.message
+import sys
 import testproto.test_pb2
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -28,4 +34,4 @@ class SeparateFileExtension(google.protobuf.message.Message):
     def HasField(self, field_name: typing.Literal["flag", b"flag"]) -> builtins.bool: ...
     def ClearField(self, field_name: typing.Literal["flag", b"flag"]) -> None: ...
 
-global___SeparateFileExtension = SeparateFileExtension
+global___SeparateFileExtension: typing_extensions.TypeAlias = SeparateFileExtension
diff --git a/test/generated/testproto/test_extensions3_pb2.pyi b/test/generated/testproto/test_extensions3_pb2.pyi
@@ -9,9 +9,15 @@ import google.protobuf.descriptor_pb2
 import google.protobuf.internal.containers
 import google.protobuf.internal.extension_dict
 import google.protobuf.message
+import sys
 import testproto.test3_pb2
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -22,7 +28,7 @@ class MessageOptionsTestMsg(google.protobuf.message.Message):
         self,
     ) -> None: ...
 
-global___MessageOptionsTestMsg = MessageOptionsTestMsg
+global___MessageOptionsTestMsg: typing_extensions.TypeAlias = MessageOptionsTestMsg
 
 TEST_FIELD_EXTENSION_FIELD_NUMBER: builtins.int
 SCALAR_OPTION_FIELD_NUMBER: builtins.int

diff --git a/test/generated/testproto/test_no_generic_services_pb2.pyi b/test/generated/testproto/test_no_generic_services_pb2.pyi
@@ -6,8 +6,14 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
+import sys
 import typing
 
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
+
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 
 @typing.final
@@ -24,4 +30,4 @@ class Simple3(google.protobuf.message.Message):
     def HasField(self, field_name: typing.Literal["a_string", b"a_string"]) -> builtins.bool: ...
     def ClearField(self, field_name: typing.Literal["a_string", b"a_string"]) -> None: ...
 
-global___Simple3 = Simple3
+global___Simple3: typing_extensions.TypeAlias = Simple3