Use typing_extension.Literal in codegen for Python 3.7

Also don't install typing-extension on Python 3.8 and above
edgedb · Oct 20, 2022 · 6d0d6ab · 6d0d6ab
1 parent 23dd42e
commit 6d0d6ab
Show file tree

Hide file tree

Showing 4 changed files with 31 additions and 20 deletions.
diff --git a/edgedb/codegen/generator.py b/edgedb/codegen/generator.py
@@ -434,26 +434,27 @@ def _generate_code(
                 print(file=buf)
                 self._imports.add("typing")
                 if SYS_VERSION_INFO >= (3, 8):
-                    for el_name, el_code in link_props:
-                        print(f"{INDENT}@typing.overload", file=buf)
-                        print(
-                            f'{INDENT}def __getitem__'
-                            f'(self, key: typing.Literal["@{el_name}"]) '
-                            f'-> {el_code}:',
-                            file=buf,
-                        )
-                        print(f"{INDENT}{INDENT}...", file=buf)
-                        print(file=buf)
+                    typing_literal = "typing.Literal"
+                else:
+                    self._imports.add("typing_extensions")
+                    typing_literal = "typing_extensions.Literal"
+                for el_name, el_code in link_props:
+                    print(f"{INDENT}@typing.overload", file=buf)
+                    print(
+                        f'{INDENT}def __getitem__'
+                        f'(self, key: {typing_literal}["@{el_name}"]) '
+                        f'-> {el_code}:',
+                        file=buf,
+                    )
+                    print(f"{INDENT}{INDENT}...", file=buf)
+                    print(file=buf)
                 print(
                     f"{INDENT}def __getitem__(self, key: str) -> typing.Any:",
                     file=buf,
                 )
-                if SYS_VERSION_INFO >= (3, 8):
-                    print(
-                        f"{INDENT}{INDENT}raise NotImplementedError", file=buf
-                    )
-                else:
-                    print(f"{INDENT}{INDENT}...", file=buf)
+                print(
+                    f"{INDENT}{INDENT}raise NotImplementedError", file=buf
+                )
 
             self._defs[rv] = buf.getvalue().strip()
 

diff --git a/edgedb/credentials.py b/edgedb/credentials.py
@@ -1,11 +1,12 @@
 import os
 import pathlib
+import sys
 import typing
 import json
 
-try:
+if sys.version_info >= (3, 8):
     from typing import TypedDict
-except ImportError:
+else:
     from typing_extensions import TypedDict
 
 from . import platform

diff --git a/setup.py b/setup.py
@@ -338,7 +338,7 @@ def finalize_options(self):
     cmdclass={'build_ext': build_ext},
     test_suite='tests.suite',
     install_requires=[
-        'typing-extensions>=3.10.0',
+        'typing-extensions>=3.10.0; python_version < "3.8.0"',
         'certifi>=2021.5.30; platform_system == "Windows"',
     ],
     extras_require=EXTRA_DEPENDENCIES,

diff --git a/tests/codegen/test-project2/object/link_prop_async_edgeql.py.assert b/tests/codegen/test-project2/object/link_prop_async_edgeql.py.assert
@@ -7,6 +7,7 @@ import dataclasses
 import datetime
 import edgedb
 import typing
+import typing_extensions
 import uuid
 
 
@@ -31,9 +32,17 @@ class LinkPropResultFriendsItem(NoPydanticValidation):
     id: uuid.UUID
     name: str
 
-    def __getitem__(self, key: str) -> typing.Any:
+    @typing.overload
+    def __getitem__(self, key: typing_extensions.Literal["@created_at"]) -> typing.Optional[datetime.datetime]:
+        ...
+
+    @typing.overload
+    def __getitem__(self, key: typing_extensions.Literal["@strength"]) -> typing.Optional[float]:
         ...
 
+    def __getitem__(self, key: str) -> typing.Any:
+        raise NotImplementedError
+
 
 async def link_prop(
     client: edgedb.AsyncIOClient,