fixer_util.py 14.2 KB
Newer Older
1 2 3
"""Utility functions, node construction macros, etc."""
# Author: Collin Winter

Benjamin Peterson's avatar
Benjamin Peterson committed
4 5
from itertools import islice

6
# Local imports
7 8 9 10
from .pgen2 import token
from .pytree import Leaf, Node
from .pygram import python_symbols as syms
from . import patcomp
11 12 13 14 15 16 17 18


###########################################################
### Common node-construction "macros"
###########################################################

def KeywordArg(keyword, value):
    return Node(syms.argument,
Benjamin Peterson's avatar
Benjamin Peterson committed
19
                [keyword, Leaf(token.EQUAL, "="), value])
20 21 22 23 24 25 26 27 28 29 30 31

def LParen():
    return Leaf(token.LPAR, "(")

def RParen():
    return Leaf(token.RPAR, ")")

def Assign(target, source):
    """Build an assignment statement"""
    if not isinstance(target, list):
        target = [target]
    if not isinstance(source, list):
32
        source.prefix = " "
33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55
        source = [source]

    return Node(syms.atom,
                target + [Leaf(token.EQUAL, "=", prefix=" ")] + source)

def Name(name, prefix=None):
    """Return a NAME leaf"""
    return Leaf(token.NAME, name, prefix=prefix)

def Attr(obj, attr):
    """A node tuple for obj.attr"""
    return [obj, Node(syms.trailer, [Dot(), attr])]

def Comma():
    """A comma leaf"""
    return Leaf(token.COMMA, ",")

def Dot():
    """A period (.) leaf"""
    return Leaf(token.DOT, ".")

def ArgList(args, lparen=LParen(), rparen=RParen()):
    """A parenthesised argument list, used by Call()"""
56 57 58 59
    node = Node(syms.trailer, [lparen.clone(), rparen.clone()])
    if args:
        node.insert_child(1, Node(syms.arglist, args))
    return node
60

61
def Call(func_name, args=None, prefix=None):
62 63 64
    """A function call"""
    node = Node(syms.power, [func_name, ArgList(args)])
    if prefix is not None:
65
        node.prefix = prefix
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80
    return node

def Newline():
    """A newline literal"""
    return Leaf(token.NEWLINE, "\n")

def BlankLine():
    """A blank line"""
    return Leaf(token.NEWLINE, "")

def Number(n, prefix=None):
    return Leaf(token.NUMBER, n, prefix=prefix)

def Subscript(index_node):
    """A numeric or string subscript"""
Benjamin Peterson's avatar
Benjamin Peterson committed
81
    return Node(syms.trailer, [Leaf(token.LBRACE, "["),
82
                               index_node,
Benjamin Peterson's avatar
Benjamin Peterson committed
83
                               Leaf(token.RBRACE, "]")])
84 85 86 87 88 89 90 91 92 93

def String(string, prefix=None):
    """A string leaf"""
    return Leaf(token.STRING, string, prefix=prefix)

def ListComp(xp, fp, it, test=None):
    """A list comprehension of the form [xp for fp in it if test].

    If test is None, the "if test" part is omitted.
    """
94 95 96
    xp.prefix = ""
    fp.prefix = " "
    it.prefix = " "
97
    for_leaf = Leaf(token.NAME, "for")
98
    for_leaf.prefix = " "
99
    in_leaf = Leaf(token.NAME, "in")
100
    in_leaf.prefix = " "
101 102
    inner_args = [for_leaf, fp, in_leaf, it]
    if test:
103
        test.prefix = " "
104
        if_leaf = Leaf(token.NAME, "if")
105
        if_leaf.prefix = " "
106 107 108 109 110 111 112
        inner_args.append(Node(syms.comp_if, [if_leaf, test]))
    inner = Node(syms.listmaker, [xp, Node(syms.comp_for, inner_args)])
    return Node(syms.atom,
                       [Leaf(token.LBRACE, "["),
                        inner,
                        Leaf(token.RBRACE, "]")])

113 114 115 116
def FromImport(package_name, name_leafs):
    """ Return an import statement in the form:
        from package import name_leafs"""
    # XXX: May not handle dotted imports properly (eg, package_name='foo.bar')
117 118 119
    #assert package_name == '.' or '.' not in package_name, "FromImport has "\
    #       "not been tested with dotted package names -- use at your own "\
    #       "peril!"
120 121 122 123 124

    for leaf in name_leafs:
        # Pull the leaves out of their old tree
        leaf.remove()

Benjamin Peterson's avatar
Benjamin Peterson committed
125
    children = [Leaf(token.NAME, "from"),
126
                Leaf(token.NAME, package_name, prefix=" "),
Benjamin Peterson's avatar
Benjamin Peterson committed
127
                Leaf(token.NAME, "import", prefix=" "),
128 129 130 131 132
                Node(syms.import_as_names, name_leafs)]
    imp = Node(syms.import_from, children)
    return imp


133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162
###########################################################
### Determine whether a node represents a given literal
###########################################################

def is_tuple(node):
    """Does the node represent a tuple literal?"""
    if isinstance(node, Node) and node.children == [LParen(), RParen()]:
        return True
    return (isinstance(node, Node)
            and len(node.children) == 3
            and isinstance(node.children[0], Leaf)
            and isinstance(node.children[1], Node)
            and isinstance(node.children[2], Leaf)
            and node.children[0].value == "("
            and node.children[2].value == ")")

def is_list(node):
    """Does the node represent a list literal?"""
    return (isinstance(node, Node)
            and len(node.children) > 1
            and isinstance(node.children[0], Leaf)
            and isinstance(node.children[-1], Leaf)
            and node.children[0].value == "["
            and node.children[-1].value == "]")


###########################################################
### Misc
###########################################################

163 164 165
def parenthesize(node):
    return Node(syms.atom, [LParen(), node, RParen()])

166 167 168 169

consuming_calls = set(["sorted", "list", "set", "any", "all", "tuple", "sum",
                       "min", "max"])

170 171
def attr_chain(obj, attr):
    """Follow an attribute chain.
172

173 174 175
    If you have a chain of objects where a.foo -> b, b.foo-> c, etc,
    use this to iterate over all objects in the chain. Iteration is
    terminated by getattr(x, attr) is None.
176

177 178 179
    Args:
        obj: the starting object
        attr: the name of the chaining attribute
180

181 182 183 184 185 186 187 188
    Yields:
        Each successive object in the chain.
    """
    next = getattr(obj, attr)
    while next:
        yield next
        next = getattr(next, attr)

189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226
p0 = """for_stmt< 'for' any 'in' node=any ':' any* >
        | comp_for< 'for' any 'in' node=any any* >
     """
p1 = """
power<
    ( 'iter' | 'list' | 'tuple' | 'sorted' | 'set' | 'sum' |
      'any' | 'all' | (any* trailer< '.' 'join' >) )
    trailer< '(' node=any ')' >
    any*
>
"""
p2 = """
power<
    'sorted'
    trailer< '(' arglist<node=any any*> ')' >
    any*
>
"""
pats_built = False
def in_special_context(node):
    """ Returns true if node is in an environment where all that is required
        of it is being itterable (ie, it doesn't matter if it returns a list
        or an itterator).
        See test_map_nochange in test_fixers.py for some examples and tests.
        """
    global p0, p1, p2, pats_built
    if not pats_built:
        p1 = patcomp.compile_pattern(p1)
        p0 = patcomp.compile_pattern(p0)
        p2 = patcomp.compile_pattern(p2)
        pats_built = True
    patterns = [p0, p1, p2]
    for pattern, parent in zip(patterns, attr_chain(node, "parent")):
        results = {}
        if pattern.match(parent, results) and results["node"] is node:
            return True
    return False

227 228 229 230
def is_probably_builtin(node):
    """
    Check that something isn't an attribute or function name etc.
    """
231
    prev = node.prev_sibling
232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249
    if prev is not None and prev.type == token.DOT:
        # Attribute lookup.
        return False
    parent = node.parent
    if parent.type in (syms.funcdef, syms.classdef):
        return False
    if parent.type == syms.expr_stmt and parent.children[0] is node:
        # Assignment.
        return False
    if parent.type == syms.parameters or \
            (parent.type == syms.typedargslist and (
            (prev is not None and prev.type == token.COMMA) or
            parent.children[0] is node
            )):
        # The name of an argument.
        return False
    return True

Benjamin Peterson's avatar
Benjamin Peterson committed
250 251 252 253 254 255 256 257 258 259
def find_indentation(node):
    """Find the indentation of *node*."""
    while node is not None:
        if node.type == syms.suite and len(node.children) > 2:
            indent = node.children[1]
            if indent.type == token.INDENT:
                return indent.value
        node = node.parent
    return ""

260 261 262 263 264 265 266 267 268 269 270 271 272
###########################################################
### The following functions are to find bindings in a suite
###########################################################

def make_suite(node):
    if node.type == syms.suite:
        return node
    node = node.clone()
    parent, node.parent = node.parent, None
    suite = Node(syms.suite, [node])
    suite.parent = parent
    return suite

273 274
def find_root(node):
    """Find the top level namespace."""
275 276 277 278 279
    # Scamper up to the top level namespace
    while node.type != syms.file_input:
        assert node.parent, "Tree is insane! root found before "\
                           "file_input node was found."
        node = node.parent
280
    return node
281

282 283 284 285 286 287
def does_tree_import(package, name, node):
    """ Returns true if name is imported from package at the
        top level of the tree which node belongs to.
        To cover the case of an import like 'import foo', use
        None for the package and 'foo' for the name. """
    binding = find_binding(name, find_root(node), package)
288 289
    return bool(binding)

290 291 292 293 294 295 296 297
def is_import(node):
    """Returns true if the node is an import statement."""
    return node.type in (syms.import_name, syms.import_from)

def touch_import(package, name, node):
    """ Works like `does_tree_import` but adds an import statement
        if it was not imported. """
    def is_import_stmt(node):
Benjamin Peterson's avatar
Benjamin Peterson committed
298 299
        return (node.type == syms.simple_stmt and node.children and
                is_import(node.children[0]))
300 301 302 303 304 305 306 307 308 309 310 311 312 313 314 315 316 317 318 319 320 321

    root = find_root(node)

    if does_tree_import(package, name, root):
        return

    # figure out where to insert the new import.  First try to find
    # the first import and then skip to the last one.
    insert_pos = offset = 0
    for idx, node in enumerate(root.children):
        if not is_import_stmt(node):
            continue
        for offset, node2 in enumerate(root.children[idx:]):
            if not is_import_stmt(node2):
                break
        insert_pos = idx + offset
        break

    # if there are no imports where we can insert, find the docstring.
    # if that also fails, we stick to the beginning of the file
    if insert_pos == 0:
        for idx, node in enumerate(root.children):
Benjamin Peterson's avatar
Benjamin Peterson committed
322 323
            if (node.type == syms.simple_stmt and node.children and
               node.children[0].type == token.STRING):
324 325 326 327 328
                insert_pos = idx + 1
                break

    if package is None:
        import_ = Node(syms.import_name, [
Benjamin Peterson's avatar
Benjamin Peterson committed
329 330
            Leaf(token.NAME, "import"),
            Leaf(token.NAME, name, prefix=" ")
331 332
        ])
    else:
Benjamin Peterson's avatar
Benjamin Peterson committed
333
        import_ = FromImport(package, [Leaf(token.NAME, name, prefix=" ")])
334 335 336 337 338

    children = [import_, Newline()]
    root.insert_child(insert_pos, Node(syms.simple_stmt, children))


339 340 341 342 343 344 345 346 347 348 349 350 351 352 353 354 355 356 357 358 359 360 361 362 363 364 365 366 367 368 369 370 371
_def_syms = set([syms.classdef, syms.funcdef])
def find_binding(name, node, package=None):
    """ Returns the node which binds variable name, otherwise None.
        If optional argument package is supplied, only imports will
        be returned.
        See test cases for examples."""
    for child in node.children:
        ret = None
        if child.type == syms.for_stmt:
            if _find(name, child.children[1]):
                return child
            n = find_binding(name, make_suite(child.children[-1]), package)
            if n: ret = n
        elif child.type in (syms.if_stmt, syms.while_stmt):
            n = find_binding(name, make_suite(child.children[-1]), package)
            if n: ret = n
        elif child.type == syms.try_stmt:
            n = find_binding(name, make_suite(child.children[2]), package)
            if n:
                ret = n
            else:
                for i, kid in enumerate(child.children[3:]):
                    if kid.type == token.COLON and kid.value == ":":
                        # i+3 is the colon, i+4 is the suite
                        n = find_binding(name, make_suite(child.children[i+4]), package)
                        if n: ret = n
        elif child.type in _def_syms and child.children[1].value == name:
            ret = child
        elif _is_import_binding(child, name, package):
            ret = child
        elif child.type == syms.simple_stmt:
            ret = find_binding(name, child, package)
        elif child.type == syms.expr_stmt:
372 373
            if _find(name, child.children[0]):
                ret = child
374 375 376 377

        if ret:
            if not package:
                return ret
378
            if is_import(ret):
379 380 381 382 383 384 385 386 387 388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413
                return ret
    return None

_block_syms = set([syms.funcdef, syms.classdef, syms.trailer])
def _find(name, node):
    nodes = [node]
    while nodes:
        node = nodes.pop()
        if node.type > 256 and node.type not in _block_syms:
            nodes.extend(node.children)
        elif node.type == token.NAME and node.value == name:
            return node
    return None

def _is_import_binding(node, name, package=None):
    """ Will reuturn node if node will import name, or node
        will import * from package.  None is returned otherwise.
        See test cases for examples. """

    if node.type == syms.import_name and not package:
        imp = node.children[1]
        if imp.type == syms.dotted_as_names:
            for child in imp.children:
                if child.type == syms.dotted_as_name:
                    if child.children[2].value == name:
                        return node
                elif child.type == token.NAME and child.value == name:
                    return node
        elif imp.type == syms.dotted_as_name:
            last = imp.children[-1]
            if last.type == token.NAME and last.value == name:
                return node
        elif imp.type == token.NAME and imp.value == name:
            return node
    elif node.type == syms.import_from:
414
        # str(...) is used to make life easier here, because
415
        # from a.b import parses to ['import', ['a', '.', 'b'], ...]
416
        if package and str(node.children[1]).strip() != package:
417 418
            return None
        n = node.children[3]
Benjamin Peterson's avatar
Benjamin Peterson committed
419
        if package and _find("as", n):
420 421 422 423 424 425 426 427 428 429 430 431 432
            # See test_from_import_as for explanation
            return None
        elif n.type == syms.import_as_names and _find(name, n):
            return node
        elif n.type == syms.import_as_name:
            child = n.children[2]
            if child.type == token.NAME and child.value == name:
                return node
        elif n.type == token.NAME and n.value == name:
            return node
        elif package and n.type == token.STAR:
            return node
    return None