Prefer typing over typing_extensions when supported (#400)

Inspired by #362

Author: nipunn1313

CHANGELOG.md

@@ -4,6 +4,7 @@
     - Allows us to use a single extension number
     - Applies to casttype,keytype,valuetype
     - Deprecate (but still support) the old-style extension
+- Prefer importing from `typing` over `typing_extensions` on new enough python versions
 - Support emitting module docstrings
 - Make generated code flake8 compatible
 - Convert extensions.proto to proto3

README.md

@@ -270,6 +270,7 @@ black .
 - [@sodul](https://github.com/sodul)
 - [@miaachan](https://github.com/miaachan)
 - [@Alphadelta14](https://github.com/Alphadelta14)
+- [@fergyfresh](https://github.com/fergyfresh)
 
 Licence etc.
 ------------

mypy_protobuf/extensions_pb2.pyi

@@ -7,7 +7,12 @@ import google.protobuf.descriptor
 import google.protobuf.descriptor_pb2
 import google.protobuf.internal.extension_dict
 import google.protobuf.message
-import typing_extensions
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

mypy_protobuf/main.py

@@ -12,6 +12,7 @@
     Iterable,
     Iterator,
     List,
+    Optional,
     Set,
     Sequence,
     Tuple,
@@ -157,6 +158,7 @@ def __init__(
         # dictionary of x->(y,z) for `from {x} import {y} as {z}`
         # if {z} is None, then it shortens to `from {x} import {y}`
         self.from_imports: Dict[str, Set[Tuple[str, str | None]]] = defaultdict(set)
+        self.typing_extensions_min: Optional[Tuple[int, int]] = None
 
         # Comments
         self.source_code_info_by_scl = {tuple(location.path): location for location in fd.source_code_info.location}
@@ -165,6 +167,16 @@ def _import(self, path: str, name: str) -> str:
         """Imports a stdlib path and returns a handle to it
         eg. self._import("typing", "Literal") -> "Literal"
         """
+        if path == "typing_extensions":
+            stabilization = {
+                "Literal": (3, 8),
+                "TypeAlias": (3, 10),
+            }
+            assert name in stabilization
+            if not self.typing_extensions_min or self.typing_extensions_min < stabilization[name]:
+                self.typing_extensions_min = stabilization[name]
+            return "typing_extensions." + name
+
         imp = path.replace("/", ".")
         if self.readable_stubs:
             self.from_imports[imp].add((name, None))
@@ -864,8 +876,16 @@ def write(self) -> str:
                 # n,n to force a reexport (from x import y as y)
                 self.from_imports[reexport_imp].update((n, n) for n in names)
 
+        if self.typing_extensions_min:
+            self.imports.add("sys")
         for pkg in sorted(self.imports):
             self._write_line(f"import {pkg}")
+        if self.typing_extensions_min:
+            self._write_line("")
+            self._write_line(f"if sys.version_info >= {self.typing_extensions_min}:")
+            self._write_line("    import typing as typing_extensions")
+            self._write_line("else:")
+            self._write_line("    import typing_extensions")
 
         for pkg, items in sorted(self.from_imports.items()):
             self._write_line(f"from {pkg} import (")

test/generated/google/protobuf/duration_pb2.pyi

@@ -35,7 +35,12 @@ import builtins
 import google.protobuf.descriptor
 import google.protobuf.internal.well_known_types
 import google.protobuf.message
-import typing_extensions
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/mypy_protobuf/extensions_pb2.pyi

@@ -7,7 +7,12 @@ import google.protobuf.descriptor
 import google.protobuf.descriptor_pb2
 import google.protobuf.internal.extension_dict
 import google.protobuf.message
-import typing_extensions
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/Capitalized/Capitalized_pb2.pyi

@@ -5,7 +5,12 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
-import typing_extensions
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/comment_special_chars_pb2.pyi

@@ -5,7 +5,12 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
-import typing_extensions
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/dot/com/test_pb2.pyi

@@ -5,7 +5,12 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
-import typing_extensions
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/grpc/dummy_pb2.pyi

@@ -5,7 +5,12 @@ https://github.com/vmagamedov/grpclib/blob/master/tests/dummy.proto"""
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
-import typing_extensions
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/inner/inner_pb2.pyi

@@ -5,8 +5,13 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
+import sys
 import testproto.test3_pb2
-import typing_extensions
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/nested/nested_pb2.pyi

@@ -6,9 +6,14 @@ import builtins
 import google.protobuf.descriptor
 import google.protobuf.internal.enum_type_wrapper
 import google.protobuf.message
+import sys
 import testproto.test3_pb2
 import typing
-import typing_extensions
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/nopackage_pb2.pyi

@@ -7,7 +7,12 @@ import collections.abc
 import google.protobuf.descriptor
 import google.protobuf.internal.containers
 import google.protobuf.message
-import typing_extensions
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/readme_enum_pb2.pyi

@@ -5,8 +5,13 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.internal.enum_type_wrapper
+import sys
 import typing
-import typing_extensions
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/test3_pb2.pyi

@@ -8,9 +8,14 @@ import google.protobuf.descriptor
 import google.protobuf.internal.containers
 import google.protobuf.internal.enum_type_wrapper
 import google.protobuf.message
+import sys
 import test.test_generated_mypy
 import typing
-import typing_extensions
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/test_extensions2_pb2.pyi

@@ -6,8 +6,13 @@ import builtins
 import google.protobuf.descriptor
 import google.protobuf.internal.extension_dict
 import google.protobuf.message
+import sys
 import testproto.test_pb2
-import typing_extensions
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/test_no_generic_services_pb2.pyi

@@ -5,7 +5,12 @@ isort:skip_file
 import builtins
 import google.protobuf.descriptor
 import google.protobuf.message
-import typing_extensions
+import sys
+
+if sys.version_info >= (3, 8):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor
 

test/generated/testproto/test_pb2.pyi

@@ -12,13 +12,18 @@ import google.protobuf.internal.enum_type_wrapper
 import google.protobuf.internal.extension_dict
 import google.protobuf.message
 import google.protobuf.service
+import sys
 import test.test_generated_mypy
 import testproto.inner.inner_pb2
 import testproto.nested.nested_pb2
 import testproto.nopackage_pb2
 import testproto.test3_pb2
 import typing
-import typing_extensions
+
+if sys.version_info >= (3, 10):
+    import typing as typing_extensions
+else:
+    import typing_extensions
 
 DESCRIPTOR: google.protobuf.descriptor.FileDescriptor