Skip to content

Commit 5faaf81

Browse files
asottileasottile
authored
Merge pull request #385 from asottile/pep604
rewrite pep604 (+ add --py310-plus)
2 parents 5e4e0dd + 6c64eaf commit 5faaf81

File tree

5 files changed

+325
-1
lines changed

5 files changed

+325
-1
lines changed

‎pyupgrade/_main.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -390,6 +390,7 @@ def _build_import_removals() -> Dict[Version, Dict[str, Tuple[str, ...]]]:
390390
((3, 7), ('generator_stop',)),
391391
((3, 8), ()),
392392
((3, 9), ()),
393+
((3, 10), ()),
393394
)
394395

395396
prev: Tuple[str, ...] = ()
@@ -874,6 +875,10 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
874875
'--py39-plus',
875876
action='store_const', dest='min_version', const=(3, 9),
876877
)
878+
parser.add_argument(
879+
'--py310-plus',
880+
action='store_const', dest='min_version', const=(3, 10),
881+
)
877882
args = parser.parse_args(argv)
878883

879884
ret = 0

‎pyupgrade/_plugins/typing_pep604.py

Lines changed: 126 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,126 @@
1+
import ast
2+
import functools
3+
import sys
4+
from typing import Iterable
5+
from typing import List
6+
from typing import Tuple
7+
8+
from tokenize_rt import Offset
9+
from tokenize_rt import Token
10+
11+
from pyupgrade._ast_helpers import ast_to_offset
12+
from pyupgrade._ast_helpers import is_name_attr
13+
from pyupgrade._data import register
14+
from pyupgrade._data import State
15+
from pyupgrade._data import TokenFunc
16+
from pyupgrade._token_helpers import CLOSING
17+
from pyupgrade._token_helpers import find_closing_bracket
18+
from pyupgrade._token_helpers import find_token
19+
from pyupgrade._token_helpers import OPENING
20+
21+
22+
def _fix_optional(i: int, tokens: List[Token]) -> None:
23+
j = find_token(tokens, i, '[')
24+
k = find_closing_bracket(tokens, j)
25+
if tokens[j].line == tokens[k].line:
26+
tokens[k] = Token('CODE', ' | None')
27+
del tokens[i:j + 1]
28+
else:
29+
tokens[j] = tokens[j]._replace(src='(')
30+
tokens[k] = tokens[k]._replace(src=')')
31+
tokens[i:j] = [Token('CODE', 'None | ')]
32+
33+
34+
def _fix_union(
35+
i: int,
36+
tokens: List[Token],
37+
*,
38+
arg: ast.expr,
39+
arg_count: int,
40+
) -> None:
41+
arg_offset = ast_to_offset(arg)
42+
j = find_token(tokens, i, '[')
43+
to_delete = []
44+
commas: List[int] = []
45+
46+
arg_depth = -1
47+
depth = 1
48+
k = j + 1
49+
while depth:
50+
if tokens[k].src in OPENING:
51+
if arg_depth == -1:
52+
to_delete.append(k)
53+
depth += 1
54+
elif tokens[k].src in CLOSING:
55+
depth -= 1
56+
if 0 < depth < arg_depth:
57+
to_delete.append(k)
58+
elif tokens[k].offset == arg_offset:
59+
arg_depth = depth
60+
elif depth == arg_depth and tokens[k].src == ',':
61+
if len(commas) >= arg_count - 1:
62+
to_delete.append(k)
63+
else:
64+
commas.append(k)
65+
66+
k += 1
67+
k -= 1
68+
69+
if tokens[j].line == tokens[k].line:
70+
del tokens[k]
71+
for comma in commas:
72+
tokens[comma] = Token('CODE', ' |')
73+
for paren in reversed(to_delete):
74+
del tokens[paren]
75+
del tokens[i:j + 1]
76+
else:
77+
tokens[j] = tokens[j]._replace(src='(')
78+
tokens[k] = tokens[k]._replace(src=')')
79+
80+
for comma in commas:
81+
tokens[comma] = Token('CODE', ' |')
82+
for paren in reversed(to_delete):
83+
del tokens[paren]
84+
del tokens[i:j]
85+
86+
87+
def _supported_version(state: State) -> bool:
88+
return (
89+
state.in_annotation and (
90+
state.settings.min_version >= (3, 10) or
91+
'annotations' in state.from_imports['__future__']
92+
)
93+
)
94+
95+
96+
@register(ast.Subscript)
97+
def visit_Subscript(
98+
state: State,
99+
node: ast.Subscript,
100+
parent: ast.AST,
101+
) -> Iterable[Tuple[Offset, TokenFunc]]:
102+
if not _supported_version(state):
103+
return
104+
105+
if is_name_attr(node.value, state.from_imports, 'typing', ('Optional',)):
106+
yield ast_to_offset(node), _fix_optional
107+
elif is_name_attr(node.value, state.from_imports, 'typing', ('Union',)):
108+
if sys.version_info >= (3, 9): # pragma: no cover (py39+)
109+
node_slice: ast.expr = node.slice
110+
elif isinstance(node.slice, ast.Index): # pragma: no cover (<py39)
111+
node_slice = node.slice.value
112+
else: # pragma: no cover (<py39)
113+
return # unexpected slice type
114+
115+
if isinstance(node_slice, ast.Tuple):
116+
if node_slice.elts:
117+
arg = node_slice.elts[0]
118+
arg_count = len(node_slice.elts)
119+
else:
120+
return # empty Union
121+
else:
122+
arg = node_slice
123+
arg_count = 1
124+
125+
func = functools.partial(_fix_union, arg=arg, arg_count=arg_count)
126+
yield ast_to_offset(node), func

‎pyupgrade/_token_helpers.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -425,7 +425,8 @@ def find_and_replace_call(
425425

426426

427427
def replace_name(i: int, tokens: List[Token], *, name: str, new: str) -> None:
428-
new_token = Token('CODE', new)
428+
# preserve token offset in case we need to match it later
429+
new_token = tokens[i]._replace(name='CODE', src=new)
429430
j = i
430431
while tokens[j].src != name:
431432
# timid: if we see a parenthesis here, skip it

‎tests/features/typing_pep604_test.py

Lines changed: 191 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,191 @@
1+
import pytest
2+
3+
from pyupgrade._data import Settings
4+
from pyupgrade._main import _fix_plugins
5+
6+
7+
@pytest.mark.parametrize(
8+
('s', 'version'),
9+
(
10+
pytest.param(
11+
'from typing import Union\n'
12+
'x: Union[int, str]\n',
13+
(3, 9),
14+
id='<3.10 Union',
15+
),
16+
pytest.param(
17+
'from typing import Optional\n'
18+
'x: Optional[str]\n',
19+
(3, 9),
20+
id='<3.10 Optional',
21+
),
22+
pytest.param(
23+
'from __future__ import annotations\n'
24+
'from typing import Union\n'
25+
'SomeAlias = Union[int, str]\n',
26+
(3, 9),
27+
id='<3.9 not in a type annotation context',
28+
),
29+
# https://github.com/python/mypy/issues/9945
30+
pytest.param(
31+
'from __future__ import annotations\n'
32+
'from typing import Union\n'
33+
'SomeAlias = Union[int, str]\n',
34+
(3, 10),
35+
id='3.10+ not in a type annotation context',
36+
),
37+
# https://github.com/python/mypy/issues/9998
38+
pytest.param(
39+
'from typing import Union\n'
40+
'def f() -> Union[()]: ...\n',
41+
(3, 10),
42+
id='3.10+ empty Union',
43+
),
44+
),
45+
)
46+
def test_fix_pep604_types_noop(s, version):
47+
assert _fix_plugins(s, settings=Settings(min_version=version)) == s
48+
49+
50+
@pytest.mark.parametrize(
51+
('s', 'expected'),
52+
(
53+
pytest.param(
54+
'from typing import Union\n'
55+
'x: Union[int, str]\n',
56+
57+
'from typing import Union\n'
58+
'x: int | str\n',
59+
60+
id='Union rewrite',
61+
),
62+
pytest.param(
63+
'x: typing.Union[int]\n',
64+
65+
'x: int\n',
66+
67+
id='Union of only one value',
68+
),
69+
pytest.param(
70+
'x: typing.Union[Foo[str, int], str]\n',
71+
72+
'x: Foo[str, int] | str\n',
73+
74+
id='Union containing a value with brackets',
75+
),
76+
pytest.param(
77+
'x: typing.Union[typing.List[str], str]\n',
78+
79+
'x: list[str] | str\n',
80+
81+
id='Union containing pep585 rewritten type',
82+
),
83+
pytest.param(
84+
'x: typing.Union[int, str,]\n',
85+
86+
'x: int | str\n',
87+
88+
id='Union trailing comma',
89+
),
90+
pytest.param(
91+
'x: typing.Union[(int, str)]\n',
92+
93+
'x: int | str\n',
94+
95+
id='Union, parenthesized tuple',
96+
),
97+
pytest.param(
98+
'x: typing.Union[\n'
99+
' int,\n'
100+
' str\n'
101+
']\n',
102+
103+
'x: (\n'
104+
' int |\n'
105+
' str\n'
106+
')\n',
107+
108+
id='Union multiple lines',
109+
),
110+
pytest.param(
111+
'x: typing.Union[\n'
112+
' int,\n'
113+
' str,\n'
114+
']\n',
115+
116+
'x: (\n'
117+
' int |\n'
118+
' str\n'
119+
')\n',
120+
121+
id='Union multiple lines with trailing commas',
122+
),
123+
pytest.param(
124+
'from typing import Optional\n'
125+
'x: Optional[str]\n',
126+
127+
'from typing import Optional\n'
128+
'x: str | None\n',
129+
130+
id='Optional rewrite',
131+
),
132+
pytest.param(
133+
'x: typing.Optional[\n'
134+
' ComplicatedLongType[int]\n'
135+
']\n',
136+
137+
'x: None | (\n'
138+
' ComplicatedLongType[int]\n'
139+
')\n',
140+
141+
id='Optional rewrite multi-line',
142+
),
143+
),
144+
)
145+
def test_fix_pep604_types(s, expected):
146+
assert _fix_plugins(s, settings=Settings(min_version=(3, 10))) == expected
147+
148+
149+
@pytest.mark.parametrize(
150+
('s', 'expected'),
151+
(
152+
pytest.param(
153+
'from __future__ import annotations\n'
154+
'from typing import Union\n'
155+
'x: Union[int, str]\n',
156+
157+
'from __future__ import annotations\n'
158+
'from typing import Union\n'
159+
'x: int | str\n',
160+
161+
id='variable annotations',
162+
),
163+
pytest.param(
164+
'from __future__ import annotations\n'
165+
'from typing import Union\n'
166+
'def f(x: Union[int, str]) -> None: ...\n',
167+
168+
'from __future__ import annotations\n'
169+
'from typing import Union\n'
170+
'def f(x: int | str) -> None: ...\n',
171+
172+
id='argument annotations',
173+
),
174+
pytest.param(
175+
'from __future__ import annotations\n'
176+
'from typing import Union\n'
177+
'def f() -> Union[int, str]: ...\n',
178+
179+
'from __future__ import annotations\n'
180+
'from typing import Union\n'
181+
'def f() -> int | str: ...\n',
182+
183+
id='return annotations',
184+
),
185+
),
186+
)
187+
def test_fix_generic_types_future_annotations(s, expected):
188+
assert _fix_plugins(s, settings=Settings(min_version=(3,))) == expected
189+
190+
191+
# TODO: test multi-line as well

‎tests/main_test.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@ def test_main_trivial():
2020
('--py37-plus',),
2121
('--py38-plus',),
2222
('--py39-plus',),
23+
('--py310-plus',),
2324
),
2425
)
2526
def test_main_noop(tmpdir, args):

0 commit comments

Comments
 (0)