Mark mangled globals as TypeAlias

nipunn1313 · Apr 24, 2024 · 4b7d1cc · 4b7d1cc
1 parent e44c1bc
commit 4b7d1cc
Showing 1 changed file with 2 additions and 2 deletions.
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: