From 9bac35526767db0761651e2075529e0d72a31cef Mon Sep 17 00:00:00 2001
From: Marc Mueller <30130371+cdce8p@users.noreply.github.com>
Date: Mon, 18 Jul 2022 22:23:13 +0200
Subject: [PATCH] Update typing rewrite with special cases

* 3.10: get_origin, get_args
* 3.11: NamedTuple, TypedDict, final, overload
---
 reorder_python_imports.py            | 13 +++++++------
 testing/generate-typing-rewrite-info | 28 +++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/reorder_python_imports.py b/reorder_python_imports.py
index 2e7d25b..33cf87b 100644
--- a/reorder_python_imports.py
+++ b/reorder_python_imports.py
@@ -421,7 +421,7 @@ def _fix_file(
 # Using:
 #     flake8-typing-imports==1.12.0
 #     mypy_extensions==0.4.3
-#     typing_extensions==4.0.1
+#     typing_extensions==4.3.0
 REPLACES[(3, 7)].update((
     'mypy_extensions=typing:NoReturn',
 ))
@@ -441,7 +441,6 @@ def _fix_file(
     'typing_extensions=typing:Text',
     'typing_extensions=typing:Type',
     'typing_extensions=typing:get_type_hints',
-    'typing_extensions=typing:overload',
 ))
 REPLACES[(3, 7)].update((
     'typing_extensions=typing:AsyncContextManager',
@@ -449,6 +448,7 @@ def _fix_file(
     'typing_extensions=typing:ChainMap',
     'typing_extensions=typing:Counter',
     'typing_extensions=typing:Deque',
+    'typing_extensions=typing:NoReturn',
 ))
 REPLACES[(3, 8)].update((
     'typing_extensions=typing:Final',
@@ -456,10 +456,6 @@ def _fix_file(
     'typing_extensions=typing:OrderedDict',
     'typing_extensions=typing:Protocol',
     'typing_extensions=typing:SupportsIndex',
-    'typing_extensions=typing:TypedDict',
-    'typing_extensions=typing:final',
-    'typing_extensions=typing:get_args',
-    'typing_extensions=typing:get_origin',
     'typing_extensions=typing:runtime_checkable',
 ))
 REPLACES[(3, 9)].update((
@@ -468,8 +464,13 @@ def _fix_file(
 REPLACES[(3, 10)].update((
     'typing_extensions=typing:Concatenate',
     'typing_extensions=typing:ParamSpec',
+    'typing_extensions=typing:ParamSpecArgs',
+    'typing_extensions=typing:ParamSpecKwargs',
     'typing_extensions=typing:TypeAlias',
     'typing_extensions=typing:TypeGuard',
+    'typing_extensions=typing:get_args',
+    'typing_extensions=typing:get_origin',
+    'typing_extensions=typing:is_typeddict',
 ))
 # END GENERATED
 
diff --git a/testing/generate-typing-rewrite-info b/testing/generate-typing-rewrite-info
index 52bbb43..dc760c6 100755
--- a/testing/generate-typing-rewrite-info
+++ b/testing/generate-typing-rewrite-info
@@ -15,6 +15,20 @@ else:
     from importlib_metadata import version
 
 
+# --- typing_extensions notes ---
+# https://github.com/python/typing_extensions#other-notes-and-limitations
+# - get_origin and get_args lack support for Annotated in Python 3.8
+#   and lack support for ParamSpecArgs and ParamSpecKwargs in 3.9.
+# - Starting with 3.11, NamedTuple and TypedDict can inherit from Generic
+# - @final was changed in Python 3.11 to set the .__final__ attribute
+# - @overload was changed in Python 3.11 to make function overloads
+#   introspectable at runtime.
+CUSTOM_TYPING_EXT_SYMBOLS = {
+    (3, 10): {'get_origin', 'get_args'},
+    (3, 11): {'NamedTuple', 'TypedDict', 'final', 'overload'},
+}
+
+
 def main() -> int:
     flake8_typing_imports_version = version('flake8-typing-imports')
     mypy_extensions_version = version('mypy_extensions')
@@ -23,7 +37,9 @@ def main() -> int:
     mypy_extensions_all = frozenset(
         a for a in dir(mypy_extensions) if a != 'Any'
     )
-    typing_extensions_all = frozenset(typing_extensions.__all__)
+    typing_extensions_all = frozenset(typing_extensions.__all__) - {
+        sym for v in CUSTOM_TYPING_EXT_SYMBOLS.values() for sym in v
+    }
 
     # some attrs are removed and then added back
     min_contiguous_versions: dict[str, flake8_typing_imports.Version] = {}
@@ -53,11 +69,13 @@ def main() -> int:
         deltas[v] = attrs - prev
         prev = attrs
 
-    mypy_extensions_added: dict[tuple[int, int], list[str]] = {}
-    typing_extensions_added: dict[tuple[int, int], list[str]] = {}
+    mypy_extensions_added: dict[tuple[int, int], set[str]] = {}
+    typing_extensions_added: dict[tuple[int, int], set[str]] = defaultdict(set)
     for v, attrs in deltas.items():
-        mypy_extensions_added[v] = sorted(attrs & mypy_extensions_all)
-        typing_extensions_added[v] = sorted(attrs & typing_extensions_all)
+        mypy_extensions_added[v] = attrs & mypy_extensions_all
+        typing_extensions_added[v] = attrs & typing_extensions_all
+    for v, attrs in CUSTOM_TYPING_EXT_SYMBOLS.items():
+        typing_extensions_added[v] |= attrs
 
     print(f'# GENERATED VIA {os.path.basename(sys.argv[0])}')
     print('# Using:')