Skip to content

add deduplication of types #1004

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

Open
walter9388 wants to merge 8 commits into
base: main
Choose a base branch
from
Open

add deduplication of types #1004

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
eca6277
initial effort
walter9388 Feb 20, 2025
dc9d5d9
clean up
walter9388 Feb 20, 2025
1fae34a
spelling
walter9388 Feb 20, 2025
90dfcab
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Feb 20, 2025
1316d66
moved modified helper function into typing_pep604.py
walter9388 Feb 22, 2025
2bbe9d7
fixed bad recursive typing xfail test
walter9388 Feb 22, 2025
2a064a8
moved other modified helper function into typing_pep604.py
walter9388 Feb 22, 2025
94931ad
fixed flake8 failures
walter9388 Feb 22, 2025
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
109 changes: 101 additions & 8 deletions pyupgrade/_plugins/typing_pep604.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,22 +14,36 @@
from pyupgrade._data import register
from pyupgrade._data import State
from pyupgrade._data import TokenFunc
from pyupgrade._token_helpers import find_closing_bracket
from pyupgrade._token_helpers import _OPENING
from pyupgrade._token_helpers import find_op
from pyupgrade._token_helpers import is_close
from pyupgrade._token_helpers import is_open


def _fix_optional(i: int, tokens: list[Token]) -> None:
j = find_op(tokens, i, '[')
k = find_closing_bracket(tokens, j)
k, contains_none = _find_closing_bracket_and_if_contains_none(tokens, j)
Copy link
Copy Markdown
Author

@walter9388 walter9388 Feb 22, 2025

Choose a reason for hiding this comment

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

Modified the general find_closing_bracket function to also check for whether the optional block already contains None in the same pass.

Copy link
Copy Markdown
Owner

@asottile asottile Apr 19, 2025

Choose a reason for hiding this comment

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

typically I add helper functions above where they're called rather than below

if tokens[j].line == tokens[k].line:
tokens[k] = Token('CODE', ' | None')
if contains_none:
del tokens[k]
else:
tokens[k:k + 1] = [
Token('UNIMPORTANT_WS', ' '),
Token('CODE', '| '),
Token('CODE', 'None'),
]
Comment on lines +30 to +34
Copy link
Copy Markdown
Author

@walter9388 walter9388 Feb 22, 2025

Choose a reason for hiding this comment

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

The reason for changing the single token containing | None to explicit whitespace, | and None is for the deduplication and whitespace removal functions used in _fix_union. This also applies to the multiline version a few lines below.

Copy link
Copy Markdown
Owner

@asottile asottile Apr 19, 2025

Choose a reason for hiding this comment

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

it might be helpful to do all the tokens then -- | would be 'OP' and None would be 'NAME'

del tokens[i:j + 1]
else:
tokens[j] = tokens[j]._replace(src='(')
tokens[k] = tokens[k]._replace(src=')')
tokens[i:j] = [Token('CODE', 'None | ')]
if contains_none:
del tokens[i:j]
else:
tokens[i:j] = [
Token('CODE', 'None'),
Token('UNIMPORTANT_WS', ' '),
Token('CODE', '| '),
]


def _fix_union(
Expand All @@ -43,6 +57,8 @@ def _fix_union(
open_parens = []
commas = []
coding_depth = None
top_level_breaks = []
lines_with_comments = []

j = find_op(tokens, i, '[')
k = j + 1
Expand Down Expand Up @@ -70,11 +86,16 @@ def _fix_union(
parens_done.append((paren_depth, (open_paren, k)))

depth -= 1
elif tokens[k].src == ',':
commas.append((depth, k))

elif tokens[k].src.strip() in [',', '|']:
if tokens[k].src.strip() == ',':
commas.append((depth, k))
if depth == 1:
top_level_breaks.append(k)
elif tokens[k].name == 'COMMENT':
lines_with_comments.append(tokens[k].line)
k += 1
k -= 1
top_level_breaks.append(k)

assert coding_depth is not None
assert not open_parens, open_parens
Expand All @@ -95,12 +116,15 @@ def _fix_union(
else:
comma_positions = []

to_delete.sort()
to_delete += _find_duplicated_types(
tokens, j, top_level_breaks, lines_with_comments,
)

if tokens[j].line == tokens[k].line:
del tokens[k]
for comma in comma_positions:
tokens[comma] = Token('CODE', ' |')
to_delete.sort()
for paren in reversed(to_delete):
del tokens[paren]
del tokens[i:j + 1]
Expand All @@ -110,11 +134,80 @@ def _fix_union(

for comma in comma_positions:
tokens[comma] = Token('CODE', ' |')
to_delete += _remove_consecutive_unimportant_ws(
tokens, [x for x in range(j, k) if x not in to_delete],
)
Comment on lines +137 to +139
Copy link
Copy Markdown
Author

@walter9388 walter9388 Feb 22, 2025

Choose a reason for hiding this comment

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

Not convinced this is the best approach to remove whitespace, but not sure about what to do in situations where lines are completely deleted other than comments. I have written a niche test for this situation in test id='duplicated types in multi-line nested unions or optionals'.

to_delete.sort()
for paren in reversed(to_delete):
del tokens[paren]
del tokens[i:j]


def _find_closing_bracket_and_if_contains_none(
tokens: list[Token],
i: int,
) -> tuple[int, bool]:
assert tokens[i].src in _OPENING
depth = 1
i += 1
contains_none = False
while depth:
if is_open(tokens[i]):
depth += 1
elif is_close(tokens[i]):
depth -= 1
elif depth == 1 and tokens[i].matches(name='NAME', src='None'):
contains_none = True
i += 1
return i - 1, contains_none


def _find_duplicated_types(
tokens: list[Token],
opening_bracket: int,
depth_1_commas: list[int],
lines_with_comments: list[int],
) -> list[int]:
unique_names = []
to_delete = []
i = opening_bracket + 1
for d1c in depth_1_commas:
important_tokens = [
x
for x in range(i, d1c)
if tokens[x].name
not in (
['COMMENT']
if tokens[x].line not in lines_with_comments
else ['COMMENT', 'NL', 'UNIMPORTANT_WS']
)
]
type_ = ''.join([tokens[k].src.lstrip() for k in important_tokens])
if type_[0] in [',', '|']:
type_ = type_[1:].lstrip()
if type_ in unique_names:
to_delete += important_tokens
else:
unique_names.append(type_)
i = d1c
return to_delete


def _remove_consecutive_unimportant_ws(
tokens: list[Token], idxs: list[int],
) -> list[int]:
to_delete = []
prev_name = ''
for kk in idxs:
Copy link
Copy Markdown
Owner

@asottile asottile Apr 19, 2025

Choose a reason for hiding this comment

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

why kk?

if prev_name == 'UNIMPORTANT_WS':
if tokens[kk].name == 'UNIMPORTANT_WS':
to_delete.append(kk)
elif tokens[kk].src == ' |':
tokens[kk] = Token('CODE', '|')
prev_name = tokens[kk].name
return to_delete


def _supported_version(state: State) -> bool:
return (
state.in_annotation and (
Expand Down
80 changes: 80 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 @@ -242,6 +242,86 @@ def f(x: int | str) -> None: ...

id='optional, 3.12: ignore close brace in fstring',
),
pytest.param(
'from typing import Optional, Union\n'
'def f(x: Optional[Union[int, None]]): pass\n'
'def g(x: Union[Optional[int], None]): pass\n'
'def h(x: Union[Union[int, None], None]): pass\n'
'def i(x: Union[int, int, None]): pass\n'
'def j(x: Union[Union[int, None], int]): pass\n'
'def k(x: Union[Union[int, None], int]): pass # comment\n'
'def l(x: Union[Union[Union[Union[a, b], c], d], a]): pass\n'
'def m(x: Union[a.b | a.c, a.b, list[str], str]): pass\n',

'from typing import Optional, Union\n'
'def f(x: int | None): pass\n'
'def g(x: int | None): pass\n'
'def h(x: int | None): pass\n'
'def i(x: int | None): pass\n'
'def j(x: int | None): pass\n'
'def k(x: int | None): pass # comment\n'
'def l(x: a | b | c | d): pass\n'
'def m(x: a.b | a.c | list[str] | str): pass\n',

id='duplicated types in nested unions or optionals',
),
pytest.param(
'from typing import Optional, Union\n'
'f: Optional[\n'
' Union[int, None]\n'
']\n'
'g: Union[\n'
' int,\n'
' int,\n'
' None,\n'
']\n'
'h: Union[\n'
' Union[int, None],\n'
' int,\n'
' None,\n'
' Optional[int],\n'
']\n'
'i: Union[\n'
' Union[int, None], # comment 1\n'
' int, # comment 2\n'
' None, # comment 3\n'
' Optional[int], # comment 4\n'
' Optional[str], # comment 5\n'
']\n',

'from typing import Optional, Union\n'
'f: (\n'
' int | None\n'
')\n'
'g: (\n'
' int |\n'
' None\n'
')\n'
'h: (\n'
' int | None\n'
')\n'
'i: (\n'
' int | None # comment 1\n'
' # comment 2\n'
' # comment 3\n'
' | # comment 4\n'
' str # comment 5\n'
')\n',

id='duplicated types in multi-line nested unions or optionals',
),
pytest.param(
'from typing import Union\n'
'def f(x: Union[list[Union[int, str]], list[Union[str, int]]]):\n'
' pass\n',

'from typing import Union\n'
'def f(x: list[int | str]):\n'
' pass\n',

id='general duplicated types',
marks=pytest.mark.xfail(reason='requires deeper type evaluation'),
),
),
)
def test_fix_pep604_types(s, expected):
Expand Down