Repository: asottile/pyupgrade

Issue:
IndentationError: unindent does not match any outer indentation level
To reproduce:
```python
# t.py

from __future__ import annotations

from typing import (
    Sequence,
    Union,
)


def sort_index(
    self,
    ascending: Union[Union[bool, int], Sequence[Union[bool, int]]],
):
    indexer = get_indexer_indexer(
        target, level, ascending, kind, na_position, sort_remaining, key
    )
```

Then:
```console
$ pyupgrade t.py 
Traceback (most recent call last):
  File "/home/marco/tmp/venv/bin/pyupgrade", line 8, in <module>
    sys.exit(main())
  File "/home/marco/tmp/venv/lib/python3.8/site-packages/pyupgrade/_main.py", line 889, in main
    ret |= _fix_file(filename, args)
  File "/home/marco/tmp/venv/lib/python3.8/site-packages/pyupgrade/_main.py", line 837, in _fix_file
    contents_text = _fix_tokens(contents_text, min_version=args.min_version)
  File "/home/marco/tmp/venv/lib/python3.8/site-packages/pyupgrade/_main.py", line 481, in _fix_tokens
    tokens = src_to_tokens(contents_text)
  File "/home/marco/tmp/venv/lib/python3.8/site-packages/tokenize_rt.py", line 68, in src_to_tokens
    for (
  File "/usr/lib/python3.8/tokenize.py", line 512, in _tokenize
    raise IndentationError(
  File "<tokenize>", line 15
    )
    ^
IndentationError: unindent does not match any outer indentation level
```

Removing either the `indexer = ` line, or the `ascending: ` annotation, makes `pyupgrade` pass. Keeping them both in, it fails.

Spotted in pandas, when commit 78371ccbd8744326e8535ac8bb9a661b073dfd88 meant that `pyupgrade` could no longer be run on `pandas/core/generic.py` without `--keep-runtime-typing`

----

version info:
```console
$ pip freeze | grep pyupgrade
pyupgrade==2.10.0
$ pip freeze | grep tokenize-rt
tokenize-rt==4.1.0
$ python --version
Python 3.8.5
```

Comments:

Patch:
diff --git a/pyupgrade/_plugins/typing_pep604.py b/pyupgrade/_plugins/typing_pep604.py
index b5e22e1..6ffcc20 100644
--- a/pyupgrade/_plugins/typing_pep604.py
+++ b/pyupgrade/_plugins/typing_pep604.py
@@ -5,6 +5,7 @@ from typing import Iterable
 from typing import List
 from typing import Tuple
 
+from tokenize_rt import NON_CODING_TOKENS
 from tokenize_rt import Offset
 from tokenize_rt import Token
 
@@ -35,40 +36,70 @@ 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
+    parens_done = []
+    open_parens = []
+    commas = []
+    coding_depth = None
+
+    j = find_token(tokens, i, '[')
     k = j + 1
     while depth:
+        # it's possible our first coding token is a close paren
+        # so make sure this is separate from the if chain below
+        if (
+                tokens[k].name not in NON_CODING_TOKENS and
+                tokens[k].src != '(' and
+                coding_depth is None
+        ):
+            if tokens[k].src == ')':  # the coding token was an empty tuple
+                coding_depth = depth - 1
+            else:
+                coding_depth = depth
+
         if tokens[k].src in OPENING:
-            if arg_depth == -1:
-                to_delete.append(k)
+            if tokens[k].src == '(':
+                open_parens.append((depth, k))
+
             depth += 1
         elif tokens[k].src in CLOSING:
+            if tokens[k].src == ')':
+                paren_depth, open_paren = open_parens.pop()
+                parens_done.append((paren_depth, (open_paren, k)))
+
             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)
+        elif tokens[k].src == ',':
+            commas.append((depth, k))
 
         k += 1
     k -= 1
 
+    assert coding_depth is not None
+    assert not open_parens, open_parens
+    comma_depth = min((depth for depth, _ in commas), default=sys.maxsize)
+    min_depth = min(comma_depth, coding_depth)
+
+    to_delete = [
+        paren
+        for depth, positions in parens_done
+        if depth < min_depth
+        for paren in positions
+    ]
+
+    if comma_depth <= coding_depth:
+        comma_positions = [k for depth, k in commas if depth == comma_depth]
+        if len(comma_positions) == arg_count:
+            to_delete.append(comma_positions.pop())
+    else:
+        comma_positions = []
+
+    to_delete.sort()
+
     if tokens[j].line == tokens[k].line:
         del tokens[k]
-        for comma in commas:
+        for comma in comma_positions:
             tokens[comma] = Token('CODE', ' |')
         for paren in reversed(to_delete):
             del tokens[paren]
@@ -77,7 +108,7 @@ def _fix_union(
         tokens[j] = tokens[j]._replace(src='(')
         tokens[k] = tokens[k]._replace(src=')')
 
-        for comma in commas:
+        for comma in comma_positions:
             tokens[comma] = Token('CODE', ' |')
         for paren in reversed(to_delete):
             del tokens[paren]
@@ -116,13 +147,11 @@ def visit_Subscript(
 
         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)
+        func = functools.partial(_fix_union, arg_count=arg_count)
         yield ast_to_offset(node), func


Test patch:
diff --git a/tests/_plugins/__init__.py b/tests/_plugins/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/tests/_plugins/typing_pep604_test.py b/tests/_plugins/typing_pep604_test.py
new file mode 100644
index 0000000..ca786b7
--- /dev/null
+++ b/tests/_plugins/typing_pep604_test.py
@@ -0,0 +1,28 @@
+import pytest
+from tokenize_rt import src_to_tokens
+from tokenize_rt import tokens_to_src
+
+from pyupgrade._plugins.typing_pep604 import _fix_union
+
+
+@pytest.mark.parametrize(
+    ('s', 'arg_count', 'expected'),
+    (
+        ('Union[a, b]', 2, 'a | b'),
+        ('Union[(a, b)]', 2, 'a | b'),
+        ('Union[(a,)]', 1, 'a'),
+        ('Union[(((a, b)))]', 2, 'a | b'),
+        pytest.param('Union[((a), b)]', 2, '(a) | b', id='wat'),
+        ('Union[(((a,), b))]', 2, '(a,) | b'),
+        ('Union[((a,), (a, b))]', 2, '(a,) | (a, b)'),
+        ('Union[((a))]', 1, 'a'),
+        ('Union[a()]', 1, 'a()'),
+        ('Union[a(b, c)]', 1, 'a(b, c)'),
+        ('Union[(a())]', 1, 'a()'),
+        ('Union[(())]', 1, '()'),
+    ),
+)
+def test_fix_union_edge_cases(s, arg_count, expected):
+    tokens = src_to_tokens(s)
+    _fix_union(0, tokens, arg_count=arg_count)
+    assert tokens_to_src(tokens) == expected
diff --git a/tests/features/typing_pep604_test.py b/tests/features/typing_pep604_test.py
index 64c8f5d..503bc2a 100644
--- a/tests/features/typing_pep604_test.py
+++ b/tests/features/typing_pep604_test.py
@@ -165,6 +165,15 @@ def f(x: int | str) -> None: ...
 
             id='Optional rewrite multi-line',
         ),
+        pytest.param(
+            'from typing import Union, Sequence\n'
+            'def f(x: Union[Union[A, B], Sequence[Union[C, D]]]): pass\n',
+
+            'from typing import Union, Sequence\n'
+            'def f(x: A | B | Sequence[C | D]): pass\n',
+
+            id='nested unions',
+        ),
     ),
 )
 def test_fix_pep604_types(s, expected):