Repository: asottile/pyupgrade

Issue:
PEP 604 rewrites
Now that mypy supports this, we should as well

notably it is these two replacements:

`Union[X, Y]` => `X | Y`
`Optional[X]` => `X | None`

Comments:

Patch:
diff --git a/pyupgrade/_main.py b/pyupgrade/_main.py
index 5e1a424..bf75a62 100644
--- a/pyupgrade/_main.py
+++ b/pyupgrade/_main.py
@@ -390,6 +390,7 @@ def _build_import_removals() -> Dict[Version, Dict[str, Tuple[str, ...]]]:
         ((3, 7), ('generator_stop',)),
         ((3, 8), ()),
         ((3, 9), ()),
+        ((3, 10), ()),
     )
 
     prev: Tuple[str, ...] = ()
@@ -874,6 +875,10 @@ def main(argv: Optional[Sequence[str]] = None) -> int:
         '--py39-plus',
         action='store_const', dest='min_version', const=(3, 9),
     )
+    parser.add_argument(
+        '--py310-plus',
+        action='store_const', dest='min_version', const=(3, 10),
+    )
     args = parser.parse_args(argv)
 
     ret = 0
diff --git a/pyupgrade/_plugins/typing_pep604.py b/pyupgrade/_plugins/typing_pep604.py
new file mode 100644
index 0000000..6dfc573
--- /dev/null
+++ b/pyupgrade/_plugins/typing_pep604.py
@@ -0,0 +1,126 @@
+import ast
+import functools
+import sys
+from typing import Iterable
+from typing import List
+from typing import Tuple
+
+from tokenize_rt import Offset
+from tokenize_rt import Token
+
+from pyupgrade._ast_helpers import ast_to_offset
+from pyupgrade._ast_helpers import is_name_attr
+from pyupgrade._data import register
+from pyupgrade._data import State
+from pyupgrade._data import TokenFunc
+from pyupgrade._token_helpers import CLOSING
+from pyupgrade._token_helpers import find_closing_bracket
+from pyupgrade._token_helpers import find_token
+from pyupgrade._token_helpers import OPENING
+
+
+def _fix_optional(i: int, tokens: List[Token]) -> None:
+    j = find_token(tokens, i, '[')
+    k = find_closing_bracket(tokens, j)
+    if tokens[j].line == tokens[k].line:
+        tokens[k] = Token('CODE', ' | None')
+        del tokens[i:j + 1]
+    else:
+        tokens[j] = tokens[j]._replace(src='(')
+        tokens[k] = tokens[k]._replace(src=')')
+        tokens[i:j] = [Token('CODE', 'None | ')]
+
+
+def _fix_union(
+        i: int,
+        tokens: List[Token],
+        *,
+        arg: ast.expr,
+        arg_count: int,
+) -> None:
+    arg_offset = ast_to_offset(arg)
+    j = find_token(tokens, i, '[')
+    to_delete = []
+    commas: List[int] = []
+
+    arg_depth = -1
+    depth = 1
+    k = j + 1
+    while depth:
+        if tokens[k].src in OPENING:
+            if arg_depth == -1:
+                to_delete.append(k)
+            depth += 1
+        elif tokens[k].src in CLOSING:
+            depth -= 1
+            if 0 < depth < arg_depth:
+                to_delete.append(k)
+        elif tokens[k].offset == arg_offset:
+            arg_depth = depth
+        elif depth == arg_depth and tokens[k].src == ',':
+            if len(commas) >= arg_count - 1:
+                to_delete.append(k)
+            else:
+                commas.append(k)
+
+        k += 1
+    k -= 1
+
+    if tokens[j].line == tokens[k].line:
+        del tokens[k]
+        for comma in commas:
+            tokens[comma] = Token('CODE', ' |')
+        for paren in reversed(to_delete):
+            del tokens[paren]
+        del tokens[i:j + 1]
+    else:
+        tokens[j] = tokens[j]._replace(src='(')
+        tokens[k] = tokens[k]._replace(src=')')
+
+        for comma in commas:
+            tokens[comma] = Token('CODE', ' |')
+        for paren in reversed(to_delete):
+            del tokens[paren]
+        del tokens[i:j]
+
+
+def _supported_version(state: State) -> bool:
+    return (
+        state.in_annotation and (
+            state.settings.min_version >= (3, 10) or
+            'annotations' in state.from_imports['__future__']
+        )
+    )
+
+
+@register(ast.Subscript)
+def visit_Subscript(
+        state: State,
+        node: ast.Subscript,
+        parent: ast.AST,
+) -> Iterable[Tuple[Offset, TokenFunc]]:
+    if not _supported_version(state):
+        return
+
+    if is_name_attr(node.value, state.from_imports, 'typing', ('Optional',)):
+        yield ast_to_offset(node), _fix_optional
+    elif is_name_attr(node.value, state.from_imports, 'typing', ('Union',)):
+        if sys.version_info >= (3, 9):  # pragma: no cover (py39+)
+            node_slice: ast.expr = node.slice
+        elif isinstance(node.slice, ast.Index):  # pragma: no cover (<py39)
+            node_slice = node.slice.value
+        else:  # pragma: no cover (<py39)
+            return  # unexpected slice type
+
+        if isinstance(node_slice, ast.Tuple):
+            if node_slice.elts:
+                arg = node_slice.elts[0]
+                arg_count = len(node_slice.elts)
+            else:
+                return  # empty Union
+        else:
+            arg = node_slice
+            arg_count = 1
+
+        func = functools.partial(_fix_union, arg=arg, arg_count=arg_count)
+        yield ast_to_offset(node), func
diff --git a/pyupgrade/_token_helpers.py b/pyupgrade/_token_helpers.py
index 3aa9e05..0ab815c 100644
--- a/pyupgrade/_token_helpers.py
+++ b/pyupgrade/_token_helpers.py
@@ -425,7 +425,8 @@ def find_and_replace_call(
 
 
 def replace_name(i: int, tokens: List[Token], *, name: str, new: str) -> None:
-    new_token = Token('CODE', new)
+    # preserve token offset in case we need to match it later
+    new_token = tokens[i]._replace(name='CODE', src=new)
     j = i
     while tokens[j].src != name:
         # timid: if we see a parenthesis here, skip it


Test patch:
diff --git a/tests/features/typing_pep604_test.py b/tests/features/typing_pep604_test.py
new file mode 100644
index 0000000..78d4ea3
--- /dev/null
+++ b/tests/features/typing_pep604_test.py
@@ -0,0 +1,191 @@
+import pytest
+
+from pyupgrade._data import Settings
+from pyupgrade._main import _fix_plugins
+
+
+@pytest.mark.parametrize(
+    ('s', 'version'),
+    (
+        pytest.param(
+            'from typing import Union\n'
+            'x: Union[int, str]\n',
+            (3, 9),
+            id='<3.10 Union',
+        ),
+        pytest.param(
+            'from typing import Optional\n'
+            'x: Optional[str]\n',
+            (3, 9),
+            id='<3.10 Optional',
+        ),
+        pytest.param(
+            'from __future__ import annotations\n'
+            'from typing import Union\n'
+            'SomeAlias = Union[int, str]\n',
+            (3, 9),
+            id='<3.9 not in a type annotation context',
+        ),
+        # https://github.com/python/mypy/issues/9945
+        pytest.param(
+            'from __future__ import annotations\n'
+            'from typing import Union\n'
+            'SomeAlias = Union[int, str]\n',
+            (3, 10),
+            id='3.10+ not in a type annotation context',
+        ),
+        # https://github.com/python/mypy/issues/9998
+        pytest.param(
+            'from typing import Union\n'
+            'def f() -> Union[()]: ...\n',
+            (3, 10),
+            id='3.10+ empty Union',
+        ),
+    ),
+)
+def test_fix_pep604_types_noop(s, version):
+    assert _fix_plugins(s, settings=Settings(min_version=version)) == s
+
+
+@pytest.mark.parametrize(
+    ('s', 'expected'),
+    (
+        pytest.param(
+            'from typing import Union\n'
+            'x: Union[int, str]\n',
+
+            'from typing import Union\n'
+            'x: int | str\n',
+
+            id='Union rewrite',
+        ),
+        pytest.param(
+            'x: typing.Union[int]\n',
+
+            'x: int\n',
+
+            id='Union of only one value',
+        ),
+        pytest.param(
+            'x: typing.Union[Foo[str, int], str]\n',
+
+            'x: Foo[str, int] | str\n',
+
+            id='Union containing a value with brackets',
+        ),
+        pytest.param(
+            'x: typing.Union[typing.List[str], str]\n',
+
+            'x: list[str] | str\n',
+
+            id='Union containing pep585 rewritten type',
+        ),
+        pytest.param(
+            'x: typing.Union[int, str,]\n',
+
+            'x: int | str\n',
+
+            id='Union trailing comma',
+        ),
+        pytest.param(
+            'x: typing.Union[(int, str)]\n',
+
+            'x: int | str\n',
+
+            id='Union, parenthesized tuple',
+        ),
+        pytest.param(
+            'x: typing.Union[\n'
+            '    int,\n'
+            '    str\n'
+            ']\n',
+
+            'x: (\n'
+            '    int |\n'
+            '    str\n'
+            ')\n',
+
+            id='Union multiple lines',
+        ),
+        pytest.param(
+            'x: typing.Union[\n'
+            '    int,\n'
+            '    str,\n'
+            ']\n',
+
+            'x: (\n'
+            '    int |\n'
+            '    str\n'
+            ')\n',
+
+            id='Union multiple lines with trailing commas',
+        ),
+        pytest.param(
+            'from typing import Optional\n'
+            'x: Optional[str]\n',
+
+            'from typing import Optional\n'
+            'x: str | None\n',
+
+            id='Optional rewrite',
+        ),
+        pytest.param(
+            'x: typing.Optional[\n'
+            '    ComplicatedLongType[int]\n'
+            ']\n',
+
+            'x: None | (\n'
+            '    ComplicatedLongType[int]\n'
+            ')\n',
+
+            id='Optional rewrite multi-line',
+        ),
+    ),
+)
+def test_fix_pep604_types(s, expected):
+    assert _fix_plugins(s, settings=Settings(min_version=(3, 10))) == expected
+
+
+@pytest.mark.parametrize(
+    ('s', 'expected'),
+    (
+        pytest.param(
+            'from __future__ import annotations\n'
+            'from typing import Union\n'
+            'x: Union[int, str]\n',
+
+            'from __future__ import annotations\n'
+            'from typing import Union\n'
+            'x: int | str\n',
+
+            id='variable annotations',
+        ),
+        pytest.param(
+            'from __future__ import annotations\n'
+            'from typing import Union\n'
+            'def f(x: Union[int, str]) -> None: ...\n',
+
+            'from __future__ import annotations\n'
+            'from typing import Union\n'
+            'def f(x: int | str) -> None: ...\n',
+
+            id='argument annotations',
+        ),
+        pytest.param(
+            'from __future__ import annotations\n'
+            'from typing import Union\n'
+            'def f() -> Union[int, str]: ...\n',
+
+            'from __future__ import annotations\n'
+            'from typing import Union\n'
+            'def f() -> int | str: ...\n',
+
+            id='return annotations',
+        ),
+    ),
+)
+def test_fix_generic_types_future_annotations(s, expected):
+    assert _fix_plugins(s, settings=Settings(min_version=(3,))) == expected
+
+
+# TODO: test multi-line as well
diff --git a/tests/main_test.py b/tests/main_test.py
index 1cc939c..da74a0c 100644
--- a/tests/main_test.py
+++ b/tests/main_test.py
@@ -20,6 +20,7 @@ def test_main_trivial():
         ('--py37-plus',),
         ('--py38-plus',),
         ('--py39-plus',),
+        ('--py310-plus',),
     ),
 )
 def test_main_noop(tmpdir, args):