Support pep585 rewriting when using alias

State.as_import probably has uses elsewhere, but this is sufficient to
start with.

We need to update some of the pep604 feature tests to ensure we actually
do the import.

Signed-off-by: Stephen Finucane <[email protected]>

Author: stephenfin

pyupgrade/_plugins/typing_pep585.py

@@ -36,7 +36,7 @@ def visit_Attribute(
     if (
             _should_rewrite(state) and
             isinstance(node.value, ast.Name) and
-            node.value.id == 'typing' and
+            state.as_imports.get(node.value.id) == 'typing' and
             node.attr in PEP585_BUILTINS
     ):
         func = functools.partial(

tests/features/typing_pep585_test.py

@@ -84,6 +84,15 @@ def f(x: list[str]) -> None: ...
 
             id='import of typing + typing.List',
         ),
+        pytest.param(
+            'import typing as ty\n'
+            'x: ty.List[int]\n',
+
+            'import typing as ty\n'
+            'x: list[int]\n',
+
+            id='aliased import of typing + typing.List',
+        ),
         pytest.param(
             'from typing import List\n'
             'SomeAlias = List[int]\n',

tests/features/typing_pep604_test.py

@@ -105,46 +105,58 @@ def f(x: int | str) -> None: ...
             id='Union rewrite',
         ),
         pytest.param(
+            'import typing\n'
             'x: typing.Union[int]\n',
 
+            'import typing\n'
             'x: int\n',
 
             id='Union of only one value',
         ),
         pytest.param(
+            'import typing\n'
             'x: typing.Union[Foo[str, int], str]\n',
 
+            'import typing\n'
             'x: Foo[str, int] | str\n',
 
             id='Union containing a value with brackets',
         ),
         pytest.param(
+            'import typing\n'
             'x: typing.Union[typing.List[str], str]\n',
 
+            'import typing\n'
             'x: list[str] | str\n',
 
             id='Union containing pep585 rewritten type',
         ),
         pytest.param(
+            'import typing\n'
             'x: typing.Union[int, str,]\n',
 
+            'import typing\n'
             'x: int | str\n',
 
             id='Union trailing comma',
         ),
         pytest.param(
+            'import typing\n'
             'x: typing.Union[(int, str)]\n',
 
+            'import typing\n'
             'x: int | str\n',
 
             id='Union, parenthesized tuple',
         ),
         pytest.param(
+            'import typing\n'
             'x: typing.Union[\n'
             '    int,\n'
             '    str\n'
             ']\n',
 
+            'import typing\n'
             'x: (\n'
             '    int |\n'
             '    str\n'
@@ -153,11 +165,13 @@ def f(x: int | str) -> None: ...
             id='Union multiple lines',
         ),
         pytest.param(
+            'import typing\n'
             'x: typing.Union[\n'
             '    int,\n'
             '    str,\n'
             ']\n',
 
+            'import typing\n'
             'x: (\n'
             '    int |\n'
             '    str\n'
@@ -175,10 +189,12 @@ def f(x: int | str) -> None: ...
             id='Optional rewrite',
         ),
         pytest.param(
+            'import typing\n'
             'x: typing.Optional[\n'
             '    ComplicatedLongType[int]\n'
             ']\n',
 
+            'import typing\n'
             'x: None | (\n'
             '    ComplicatedLongType[int]\n'
             ')\n',