Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

pep585 import aliases #983

Closed
wants to merge 2 commits into from
Closed

pep585 import aliases #983

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e4f87de
Track import aliases via State
stephenfin Oct 24, 2024
3bc5f61
Support pep585 rewriting when using alias
stephenfin Oct 24, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 11 additions & 2 deletions pyupgrade/_data.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,14 +27,15 @@ class Settings(NamedTuple):
class State(NamedTuple):
settings: Settings
from_imports: dict[str, set[str]]
as_imports: dict[str, str]
in_annotation: bool = False


AST_T = TypeVar('AST_T', bound=ast.AST)
TokenFunc = Callable[[int, list[Token]], None]
ASTFunc = Callable[[State, AST_T, ast.AST], Iterable[tuple[Offset, TokenFunc]]]

RECORD_FROM_IMPORTS = frozenset((
RECORDED_IMPORTS = frozenset((
'__future__',
'asyncio',
'collections',
Expand Down Expand Up @@ -75,6 +76,7 @@ def visit(
initial_state = State(
settings=settings,
from_imports=collections.defaultdict(set),
as_imports=collections.defaultdict(),
)

nodes: list[tuple[State, ast.AST, ast.AST]] = [(initial_state, tree, tree)]
Expand All @@ -91,11 +93,18 @@ def visit(
if (
isinstance(node, ast.ImportFrom) and
not node.level and
node.module in RECORD_FROM_IMPORTS
node.module in RECORDED_IMPORTS
):
state.from_imports[node.module].update(
name.name for name in node.names if not name.asname
)
elif (
isinstance(node, ast.Import)
):
state.as_imports.update({
x.asname or x.name: x.name for x in node.names
if x.name in RECORDED_IMPORTS
})

for name in reversed(node._fields):
value = getattr(node, name)
Expand Down
2 changes: 1 addition & 1 deletion pyupgrade/_plugins/typing_pep585.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Author

@stephenfin stephenfin Oct 24, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that this breaks cases like we see in typing_pep604_test.py below, where we use typing without first importing it. I think that's okay since that code is already broken in real-life, but if we wanted to avoid making those changes we could trivially change this to:

Suggested change
state.as_imports.get(node.value.id) == 'typing' and
state.as_imports.get(node.value.id, node.value.id) == 'typing' and

node.attr in PEP585_BUILTINS
):
func = functools.partial(
Expand Down
9 changes: 9 additions & 0 deletions tests/features/typing_pep585_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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',
Expand Down
16 changes: 16 additions & 0 deletions tests/features/typing_pep604_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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'
Expand All @@ -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'
Expand All @@ -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',
Expand Down
Loading