Repository: asottile/pyupgrade

Issue:
Rewrite `if PY2:` blocks in `--py3-plus`
Ideally handle these different cases:

```python
# input
from six import PY2
if PY2:
   x = 5
else:
   x = 6
# output
x = 6
```

But also for

```python
sys.version_info[0] == 2
sys.version_info < (3,)
```
etc.

- [x] `six.PY2`
- [x] `six.PY3`
- [x] `not six.PY2`
- [x] `not six.PY3`
- [ ] `sys.version_info >= (3,)`
- [ ] `sys.version_info < (3,)`
- [ ] `sys.version_info[0] == 2`
- [ ] `sys.version_info[0] == 3`

Discussion:

Patch:
diff --git a/pyupgrade.py b/pyupgrade.py
index ac8bd70..b26904c 100644
--- a/pyupgrade.py
+++ b/pyupgrade.py
@@ -961,6 +961,10 @@ class FindPy3Plus(ast.NodeVisitor):
         self.bases_to_remove = set()
 
         self.encode_calls = {}
+
+        self.if_py2_blocks = set()
+        self.if_py3_blocks = set()
+
         self.native_literals = set()
 
         self._six_from_imports = set()
@@ -1106,11 +1110,162 @@ class FindPy3Plus(ast.NodeVisitor):
 
         self.generic_visit(node)
 
-    def generic_visit(self, node):
+    def visit_If(self, node):  # type: (ast.If) -> None
+        if node.orelse and not isinstance(node.orelse[0], ast.If):
+            if (
+                    # if six.PY2:
+                    self._is_six(node.test, ('PY2',)) or
+                    # if not six.PY3:
+                    (
+                        isinstance(node.test, ast.UnaryOp) and
+                        isinstance(node.test.op, ast.Not) and
+                        self._is_six(node.test.operand, ('PY3',))
+                    )
+            ):
+                self.if_py2_blocks.add(_ast_to_offset(node))
+            elif (
+                    # if six.PY3:
+                    self._is_six(node.test, 'PY3') or
+                    # if not six.PY2:
+                    (
+                        isinstance(node.test, ast.UnaryOp) and
+                        isinstance(node.test.op, ast.Not) and
+                        self._is_six(node.test.operand, ('PY2',))
+                    )
+            ):
+                self.if_py3_blocks.add(_ast_to_offset(node))
+        self.generic_visit(node)
+
+    def generic_visit(self, node):  # type: (ast.AST) -> None
         self._previous_node = node
         super(FindPy3Plus, self).generic_visit(node)
 
 
+def _fixup_dedent_tokens(tokens):
+    """For whatever reason the DEDENT / UNIMPORTANT_WS tokens are misordered
+
+    | if True:
+    |     if True:
+    |         pass
+    |     else:
+    |^    ^- DEDENT
+    |+----UNIMPORTANT_WS
+    """
+    for i, token in enumerate(tokens):
+        if token.name == UNIMPORTANT_WS and tokens[i + 1].name == 'DEDENT':
+            tokens[i], tokens[i + 1] = tokens[i + 1], tokens[i]
+
+
+def _find_block_start(tokens, i):
+    depth = 0
+    while depth or tokens[i].src != ':':
+        depth += {'(': 1, ')': -1}.get(tokens[i].src, 0)
+        i += 1
+    return i
+
+
+class Block(
+        collections.namedtuple('Block', ('start', 'block', 'end', 'line')),
+):
+    __slots__ = ()
+
+    def _initial_indent(self, tokens):
+        if tokens[self.start].src.isspace():
+            return len(tokens[self.start].src)
+        else:
+            return 0
+
+    def _minimum_indent(self, tokens):
+        block_indent = None
+        for i in range(self.block, self.end):
+            if (
+                    tokens[i - 1].name in ('NL', 'NEWLINE') and
+                    tokens[i].name in ('INDENT', UNIMPORTANT_WS)
+            ):
+                token_indent = len(tokens[i].src)
+                if block_indent is None:
+                    block_indent = token_indent
+                else:
+                    block_indent = min(block_indent, token_indent)
+        return block_indent
+
+    def dedent(self, tokens):
+        if self.line:
+            return
+        diff = self._minimum_indent(tokens) - self._initial_indent(tokens)
+        for i in range(self.block, self.end):
+            if (
+                    tokens[i - 1].name in ('NL', 'NEWLINE') and
+                    tokens[i].name in ('INDENT', UNIMPORTANT_WS)
+            ):
+                tokens[i] = tokens[i]._replace(src=tokens[i].src[diff:])
+
+    def _trim_end(self, tokens):
+        """the tokenizer reports the end of the block at the beginning of
+        the next block
+        """
+        i = last_token = self.end - 1
+        while tokens[i].name in NON_CODING_TOKENS | {'DEDENT', 'NEWLINE'}:
+            # if we find an indented comment inside our block, keep it
+            if (
+                    tokens[i].name in {'NL', 'NEWLINE'} and
+                    tokens[i + 1].name == UNIMPORTANT_WS and
+                    len(tokens[i + 1].src) > self._initial_indent(tokens)
+            ):
+                break
+            # otherwise we've found another line to remove
+            elif tokens[i].name in {'NL', 'NEWLINE'}:
+                last_token = i
+            i -= 1
+        return self._replace(end=last_token + 1)
+
+    @classmethod
+    def find(cls, tokens, i, trim_end=False):
+        if i > 0 and tokens[i - 1].name in {'INDENT', UNIMPORTANT_WS}:
+            i -= 1
+        start = i
+        colon = _find_block_start(tokens, i)
+
+        j = colon + 1
+        while (
+                tokens[j].name != 'NEWLINE' and
+                tokens[j].name in NON_CODING_TOKENS
+        ):
+            j += 1
+
+        if tokens[j].name == 'NEWLINE':  # multi line block
+            block = j + 1
+            while tokens[j].name != 'INDENT':
+                j += 1
+            level = 1
+            j += 1
+            while level:
+                level += {'INDENT': 1, 'DEDENT': -1}.get(tokens[j].name, 0)
+                j += 1
+            if trim_end:
+                return cls(start, block, j, line=False)._trim_end(tokens)
+            else:
+                return cls(start, block, j, line=False)
+        else:  # single line block
+            block = j
+            # search forward until the NEWLINE token
+            while tokens[j].name not in {'NEWLINE', 'ENDMARKER'}:
+                j += 1
+            # we also want to include the newline in the block
+            if tokens[j].name == 'NEWLINE':  # pragma: no branch (PY2 only)
+                j += 1
+            return cls(start, block, j, line=True)
+
+
+def _find_if_else_block(tokens, i):
+    if_block = Block.find(tokens, i)
+    i = if_block.end
+    while tokens[i].src != 'else':
+        i += 1
+    else_block = Block.find(tokens, i, trim_end=True)
+    return if_block, else_block
+
+
 def _remove_decorator(tokens, i):
     while tokens[i - 1].src != '@':
         i -= 1
@@ -1236,6 +1391,8 @@ def _fix_py3_plus(contents_text):
     if not any((
             visitor.bases_to_remove,
             visitor.encode_calls,
+            visitor.if_py2_blocks,
+            visitor.if_py3_blocks,
             visitor.native_literals,
             visitor.six_add_metaclass,
             visitor.six_b,
@@ -1254,6 +1411,8 @@ def _fix_py3_plus(contents_text):
     except tokenize.TokenError:  # pragma: no cover (bpo-2180)
         return contents_text
 
+    _fixup_dedent_tokens(tokens)
+
     def _replace(i, mapping, node):
         new_token = Token('CODE', _get_tmpl(mapping, node))
         if isinstance(node, ast.Name):
@@ -1269,6 +1428,21 @@ def _fix_py3_plus(contents_text):
             continue
         elif token.offset in visitor.bases_to_remove:
             _remove_base_class(tokens, i)
+        elif token.offset in visitor.if_py2_blocks:
+            if tokens[i].src != 'if':
+                continue
+            if_block, else_block = _find_if_else_block(tokens, i)
+
+            else_block.dedent(tokens)
+            del tokens[if_block.start:else_block.block]
+        elif token.offset in visitor.if_py3_blocks:
+            if tokens[i].src != 'if':
+                continue
+            if_block, else_block = _find_if_else_block(tokens, i)
+
+            if_block.dedent(tokens)
+            del tokens[if_block.end:else_block.end]
+            del tokens[if_block.start:if_block.block]
         elif token.offset in visitor.six_type_ctx:
             _replace(i, SIX_TYPE_CTX_ATTRS, visitor.six_type_ctx[token.offset])
         elif token.offset in visitor.six_simple:
@@ -1305,16 +1479,13 @@ def _fix_py3_plus(contents_text):
             j = i + 1
             while tokens[j].src != 'class':
                 j += 1
+            class_token = j
             # then search forward for a `:` token, not inside a brace
-            depth = 0
+            j = _find_block_start(tokens, j)
             last_paren = -1
-            while depth or tokens[j].src != ':':
-                if tokens[j].src == '(':
-                    depth += 1
-                elif tokens[j].src == ')':
-                    depth -= 1
-                    last_paren = j
-                j += 1
+            for k in range(class_token, j):
+                if tokens[k].src == ')':
+                    last_paren = k
 
             if last_paren == -1:
                 tokens.insert(j, Token('CODE', '({})'.format(metaclass)))


Test patch:
diff --git a/tests/pyupgrade_test.py b/tests/pyupgrade_test.py
index 87ac751..070c9fe 100644
--- a/tests/pyupgrade_test.py
+++ b/tests/pyupgrade_test.py
@@ -1493,6 +1493,274 @@ def test_fix_encode_noop(s):
     assert _fix_py3_plus(s) == s
 
 
+@pytest.mark.parametrize(
+    's',
+    (
+        # we timidly skip `if` without `else` as it could cause a SyntaxError
+        'if six.PY2:\n'
+        '    pass',
+        # here's the case where it causes a SyntaxError
+        'if True:\n'
+        '    if six.PY2:\n'
+        '        pass\n',
+        # for now we don't attempt to rewrite `elif`
+        'if False:\n'
+        '    pass\n'
+        'elif six.PY3:\n'
+        '    pass\n'
+        'else:\n'
+        '    pass\n',
+        'if False:\n'
+        '    pass\n'
+        'elif six.PY2:\n'
+        '    pass\n'
+        'else:\n'
+        '    pass\n',
+    ),
+)
+def test_fix_py2_block_noop(s):
+    assert _fix_py3_plus(s) == s
+
+
+@pytest.mark.parametrize(
+    ('s', 'expected'),
+    (
+        pytest.param(
+            'if six.PY2:\n'
+            '    print("py2")\n'
+            'else:\n'
+            '    print("py3")\n',
+
+            'print("py3")\n',
+
+            id='six.PY2',
+        ),
+        pytest.param(
+            'if six.PY2:\n'
+            '    if True:\n'
+            '        print("py2!")\n'
+            '    else:\n'
+            '        print("???")\n'
+            'else:\n'
+            '    print("py3")\n',
+
+            'print("py3")\n',
+
+            id='six.PY2, nested ifs',
+        ),
+        pytest.param(
+            'if six.PY2: print("PY2!")\n'
+            'else:       print("PY3!")\n',
+
+            'print("PY3!")\n',
+
+            id='six.PY2, oneline',
+        ),
+        pytest.param(
+            'if six.PY2: print("PY2!")\n'
+            'else:       print("PY3!")',
+
+            'print("PY3!")',
+
+            id='six.PY2, oneline, no newline at end of file',
+        ),
+        pytest.param(
+            'if True:\n'
+            '    if six.PY2:\n'
+            '        print("PY2")\n'
+            '    else:\n'
+            '        print("PY3")\n',
+
+            'if True:\n'
+            '    print("PY3")\n',
+
+            id='six.PY2, indented',
+        ),
+        pytest.param(
+            'if six.PY2: print(1 if True else 3)\n'
+            'else:\n'
+            '   print("py3")\n',
+
+            'print("py3")\n',
+
+            id='six.PY2, `else` token inside oneline if\n',
+        ),
+        pytest.param(
+            'if six.PY2:\n'
+            '    def f():\n'
+            '        print("py2")\n'
+            'else:\n'
+            '    def f():\n'
+            '        print("py3")\n',
+
+            'def f():\n'
+            '    print("py3")\n',
+
+            id='six.PY2, multiple indents in block',
+        ),
+        pytest.param(
+            'if not six.PY2:\n'
+            '    print("py3")\n'
+            'else:\n'
+            '    print("py2")\n'
+            '\n'
+            '\n'
+            'x = 1\n',
+
+            'print("py3")\n'
+            '\n'
+            '\n'
+            'x = 1\n',
+
+            id='not six.PY2, remove second block',
+        ),
+        pytest.param(
+            'if not six.PY2:\n'
+            '    print("py3")\n'
+            'else:\n'
+            '    print("py2")',
+
+            'print("py3")\n',
+
+            id='not six.PY2, no end of line',
+        ),
+        pytest.param(
+            'if not six.PY2:\n'
+            '    print("py3")\n'
+            'else:\n'
+            '    print("py2")\n'
+            '    # ohai\n'
+            '\n'
+            'x = 1\n',
+
+            'print("py3")\n'
+            '\n'
+            'x = 1\n',
+
+            id='not six.PY2: else block ends in comment',
+        ),
+        pytest.param(
+            'if not six.PY2: print("py3")\n'
+            'else: print("py2")\n',
+
+            'print("py3")\n',
+
+            id='not six.PY2, else is single line',
+        ),
+        pytest.param(
+            'if six.PY3:\n'
+            '    print("py3")\n'
+            'else:\n'
+            '    print("py2")\n',
+
+            'print("py3")\n',
+
+            id='six.PY3',
+        ),
+        pytest.param(
+            'if True:\n'
+            '    if six.PY3:\n'
+            '        print("py3")\n'
+            '    else:\n'
+            '        print("py2")\n',
+
+            'if True:\n'
+            '    print("py3")\n',
+
+            id='indented six.PY3',
+        ),
+        pytest.param(
+            'from six import PY3\n'
+            'if not PY3:\n'
+            '    print("py2")\n'
+            'else:\n'
+            '    print("py3")\n',
+
+            'from six import PY3\n'
+            'print("py3")\n',
+
+            id='not PY3',
+        ),
+        pytest.param(
+            'def f():\n'
+            '    if six.PY2:\n'
+            '        try:\n'
+            '            yield\n'
+            '        finally:\n'
+            '            pass\n'
+            '    else:\n'
+            '        yield\n',
+
+            'def f():\n'
+            '    yield\n',
+
+            id='six.PY2, finally',
+        ),
+        pytest.param(
+            'class C:\n'
+            '    def g():\n'
+            '        pass\n'
+            '\n'
+            '    if six.PY2:\n'
+            '        def f(py2):\n'
+            '            pass\n'
+            '    else:\n'
+            '        def f(py3):\n'
+            '            pass\n'
+            '\n'
+            '    def h():\n'
+            '        pass\n',
+
+            'class C:\n'
+            '    def g():\n'
+            '        pass\n'
+            '\n'
+            '    def f(py3):\n'
+            '        pass\n'
+            '\n'
+            '    def h():\n'
+            '        pass\n',
+
+            id='six.PY2 in class\n',
+        ),
+        pytest.param(
+            'if True:\n'
+            '    if six.PY2:\n'
+            '        2\n'
+            '    else:\n'
+            '        3\n'
+            '\n'
+            '    # comment\n',
+
+            'if True:\n'
+            '    3\n'
+            '\n'
+            '    # comment\n',
+
+            id='six.PY2, comment after',
+        ),
+        pytest.param(
+            'if True:\n'
+            '    if six.PY3:\n'
+            '        3\n'
+            '    else:\n'
+            '        2\n'
+            '\n'
+            '    # comment\n',
+
+            'if True:\n'
+            '    3\n'
+            '\n'
+            '    # comment\n',
+
+            id='six.PY3, comment after',
+        ),
+    ),
+)
+def test_fix_py2_blocks(s, expected):
+    assert _fix_py3_plus(s) == expected
+
+
 @pytest.mark.parametrize(
     's',
     (