1 # Copyright 2007 Google, Inc. All Rights Reserved.

     2 # Licensed to PSF under a Contributor Agreement.

     3 

     4 """Fixer that changes xrange(...) into range(...)."""

     5 

     6 # Local imports

     7 from .. import fixer_base

     8 from ..fixer_util import Name, Call, consuming_calls

     9 from .. import patcomp

    10 

    11 

    12 class FixXrange(fixer_base.BaseFix):

    13 

    14     PATTERN = """

    15               power< (name='range'|name='xrange') trailer< '(' [any] ')' > any* >

    16               """

    17 

    18     def transform(self, node, results):

    19         name = results["name"]

    20         if name.value == "xrange":

    21             return self.transform_xrange(node, results)

    22         elif name.value == "range":

    23             return self.transform_range(node, results)

    24         else:

    25             raise ValueError(repr(name))

    26 

    27     def transform_xrange(self, node, results):

    28         name = results["name"]

    29         name.replace(Name("range", prefix=name.get_prefix()))

    30 

    31     def transform_range(self, node, results):

    32         if not self.in_special_context(node):

    33             arg = node.clone()

    34             arg.set_prefix("")

    35             call = Call(Name("list"), [arg])

    36             call.set_prefix(node.get_prefix())

    37             return call

    38         return node

    39 

    40     P1 = "power< func=NAME trailer< '(' node=any ')' > any* >"

    41     p1 = patcomp.compile_pattern(P1)

    42 

    43     P2 = """for_stmt< 'for' any 'in' node=any ':' any* >

    44             | comp_for< 'for' any 'in' node=any any* >

    45             | comparison< any 'in' node=any any*>

    46          """

    47     p2 = patcomp.compile_pattern(P2)

    48 

    49     def in_special_context(self, node):

    50         if node.parent is None:

    51             return False

    52         results = {}

    53         if (node.parent.parent is not None and

    54                self.p1.match(node.parent.parent, results) and

    55                results["node"] is node):

    56             # list(d.keys()) -> list(d.keys()), etc.

    57             return results["func"].value in consuming_calls

    58         # for ... in d.iterkeys() -> for ... in d.keys(), etc.

    59         return self.p2.match(node.parent, results) and results["node"] is node