versioned_branches.py

import ast
from typing import cast
from typing import Iterable
from typing import List
from typing import Tuple
from typing import Type
from typing import Union

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 Block


def _find_if_else_block(tokens: List[Token], i: int) -> Tuple[Block, Block]:
    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 _find_elif(tokens: List[Token], i: int) -> int:
    while tokens[i].src != 'elif':  # pragma: no cover (only for <3.8.1)
        i -= 1
    return i


def _fix_py3_block(i: int, tokens: List[Token]) -> None:
    if tokens[i].src == 'if':
        if_block = Block.find(tokens, i)
        if_block.dedent(tokens)
        del tokens[if_block.start:if_block.block]
    else:
        if_block = Block.find(tokens, _find_elif(tokens, i))
        if_block.replace_condition(tokens, [Token('NAME', 'else')])


def _fix_py2_block(i: int, tokens: List[Token]) -> None:
    if tokens[i].src == 'if':
        if_block, else_block = _find_if_else_block(tokens, i)
        else_block.dedent(tokens)
        del tokens[if_block.start:else_block.block]
    else:
        j = _find_elif(tokens, i)
        if_block, else_block = _find_if_else_block(tokens, j)
        del tokens[if_block.start:else_block.start]


def _fix_py3_block_else(i: int, tokens: List[Token]) -> None:
    if tokens[i].src == 'if':
        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]
    else:
        j = _find_elif(tokens, i)
        if_block, else_block = _find_if_else_block(tokens, j)
        del tokens[if_block.end:else_block.end]
        if_block.replace_condition(tokens, [Token('NAME', 'else')])


def _eq(test: ast.Compare, n: int) -> bool:
    return (
        isinstance(test.ops[0], ast.Eq) and
        isinstance(test.comparators[0], ast.Num) and
        test.comparators[0].n == n
    )


def _compare_to_3(
    test: ast.Compare,
    op: Union[Type[ast.cmpop], Tuple[Type[ast.cmpop], ...]],
    minor: int = 0,
) -> bool:
    min_len = 2 if minor else 1
    if not (
            isinstance(test.ops[0], op) and
            isinstance(test.comparators[0], ast.Tuple) and
            len(test.comparators[0].elts) >= min_len and
            all(isinstance(n, ast.Num) for n in test.comparators[0].elts)
    ):
        return False

    # checked above but mypy needs help
    elts = cast('List[ast.Num]', test.comparators[0].elts)

    retv = elts[0].n == 3
    offset = 1
    if minor:
        retv &= elts[1].n == minor
        offset += 1
    retv &= all(n.n == 0 for n in elts[offset:])
    return retv


@register(ast.If)
def visit_If(
        state: State,
        node: ast.If,
        parent: ast.AST,
) -> Iterable[Tuple[Offset, TokenFunc]]:
    if state.settings.min_version >= (3,) and (
        len(state.settings.min_version) >= 2
    ):
        py3_minor = state.settings.min_version[1]

    if (
            state.settings.min_version >= (3,) and (
                # if six.PY2:
                is_name_attr(node.test, state.from_imports, 'six', ('PY2',)) or
                # if not six.PY3:
                (
                    isinstance(node.test, ast.UnaryOp) and
                    isinstance(node.test.op, ast.Not) and
                    is_name_attr(
                        node.test.operand,
                        state.from_imports,
                        'six',
                        ('PY3',),
                    )
                ) or
                # sys.version_info == 2 or < (3,)
                (
                    isinstance(node.test, ast.Compare) and
                    is_name_attr(
                        node.test.left,
                        state.from_imports,
                        'sys',
                        ('version_info',),
                    ) and
                    len(node.test.ops) == 1 and (
                        _eq(node.test, 2) or
                        _compare_to_3(node.test, ast.Lt)
                    )
                ) or
                # sys.version_info < (3, n) (with n<=m)
                # or sys.version_info <= (3, n) (with n<m)
                (
                    len(state.settings.min_version) >= 2 and (
                        isinstance(node.test, ast.Compare) and
                        is_name_attr(
                            node.test.left,
                            state.from_imports,
                            'sys',
                            ('version_info',),
                        ) and
                        len(node.test.ops) == 1 and (
                            _compare_to_3(
                                node.test,
                                ast.Lt,
                                minor=py3_minor,
                            ) or any(
                                _compare_to_3(
                                    node.test,
                                    (ast.Lt, ast.LtE),
                                    minor=minor,
                                )
                                for minor in range(py3_minor)
                            )
                        )
                    )
                )
            )
    ):
        if node.orelse and not isinstance(node.orelse[0], ast.If):
            yield ast_to_offset(node), _fix_py2_block
    elif (
            state.settings.min_version >= (3,) and (
                # if six.PY3:
                is_name_attr(node.test, state.from_imports, 'six', ('PY3',)) or
                # if not six.PY2:
                (
                    isinstance(node.test, ast.UnaryOp) and
                    isinstance(node.test.op, ast.Not) and
                    is_name_attr(
                        node.test.operand,
                        state.from_imports,
                        'six',
                        ('PY2',),
                    )
                ) or
                # sys.version_info == 3 or >= (3,) or > (3,)
                (
                    isinstance(node.test, ast.Compare) and
                    is_name_attr(
                        node.test.left,
                        state.from_imports,
                        'sys',
                        ('version_info',),
                    ) and
                    len(node.test.ops) == 1 and (
                        _eq(node.test, 3) or
                        _compare_to_3(node.test, (ast.Gt, ast.GtE))
                    )
                ) or
                # sys.version_info >= (3, n) (with n<=m)
                # or sys.version_info > (3, n) (with n<m)
                (
                    len(state.settings.min_version) >= 2 and (

                        isinstance(node.test, ast.Compare) and
                        is_name_attr(
                            node.test.left,
                            state.from_imports,
                            'sys',
                            ('version_info',),
                        ) and
                        len(node.test.ops) == 1 and
                        (
                            _compare_to_3(
                                node.test,
                                ast.GtE,
                                minor=py3_minor,
                            ) or any(
                                _compare_to_3(
                                    node.test,
                                    (ast.Gt, ast.GtE),
                                    minor=minor,
                                )
                                for minor in range(py3_minor)
                            )
                        )

                    )
                )
            )
    ):
        if node.orelse and not isinstance(node.orelse[0], ast.If):
            yield ast_to_offset(node), _fix_py3_block_else
        elif not node.orelse:
            yield ast_to_offset(node), _fix_py3_block