1 """Utility functions, node construction macros, etc."""

     2 # Author: Collin Winter

     3 

     4 # Local imports

     5 from .pgen2 import token

     6 from .pytree import Leaf, Node

     7 from .pygram import python_symbols as syms

     8 from . import patcomp

     9 

    10 

    11 ###########################################################

    12 ### Common node-construction "macros"

    13 ###########################################################

    14 

    15 def KeywordArg(keyword, value):

    16     return Node(syms.argument,

    17                 [keyword, Leaf(token.EQUAL, '='), value])

    18 

    19 def LParen():

    20     return Leaf(token.LPAR, "(")

    21 

    22 def RParen():

    23     return Leaf(token.RPAR, ")")

    24 

    25 def Assign(target, source):

    26     """Build an assignment statement"""

    27     if not isinstance(target, list):

    28         target = [target]

    29     if not isinstance(source, list):

    30         source.set_prefix(" ")

    31         source = [source]

    32 

    33     return Node(syms.atom,

    34                 target + [Leaf(token.EQUAL, "=", prefix=" ")] + source)

    35 

    36 def Name(name, prefix=None):

    37     """Return a NAME leaf"""

    38     return Leaf(token.NAME, name, prefix=prefix)

    39 

    40 def Attr(obj, attr):

    41     """A node tuple for obj.attr"""

    42     return [obj, Node(syms.trailer, [Dot(), attr])]

    43 

    44 def Comma():

    45     """A comma leaf"""

    46     return Leaf(token.COMMA, ",")

    47 

    48 def Dot():

    49     """A period (.) leaf"""

    50     return Leaf(token.DOT, ".")

    51 

    52 def ArgList(args, lparen=LParen(), rparen=RParen()):

    53     """A parenthesised argument list, used by Call()"""

    54     node = Node(syms.trailer, [lparen.clone(), rparen.clone()])

    55     if args:

    56         node.insert_child(1, Node(syms.arglist, args))

    57     return node

    58 

    59 def Call(func_name, args=None, prefix=None):

    60     """A function call"""

    61     node = Node(syms.power, [func_name, ArgList(args)])

    62     if prefix is not None:

    63         node.set_prefix(prefix)

    64     return node

    65 

    66 def Newline():

    67     """A newline literal"""

    68     return Leaf(token.NEWLINE, "\n")

    69 

    70 def BlankLine():

    71     """A blank line"""

    72     return Leaf(token.NEWLINE, "")

    73 

    74 def Number(n, prefix=None):

    75     return Leaf(token.NUMBER, n, prefix=prefix)

    76 

    77 def Subscript(index_node):

    78     """A numeric or string subscript"""

    79     return Node(syms.trailer, [Leaf(token.LBRACE, '['),

    80                                index_node,

    81                                Leaf(token.RBRACE, ']')])

    82 

    83 def String(string, prefix=None):

    84     """A string leaf"""

    85     return Leaf(token.STRING, string, prefix=prefix)

    86 

    87 def ListComp(xp, fp, it, test=None):

    88     """A list comprehension of the form [xp for fp in it if test].

    89 

    90     If test is None, the "if test" part is omitted.

    91     """

    92     xp.set_prefix("")

    93     fp.set_prefix(" ")

    94     it.set_prefix(" ")

    95     for_leaf = Leaf(token.NAME, "for")

    96     for_leaf.set_prefix(" ")

    97     in_leaf = Leaf(token.NAME, "in")

    98     in_leaf.set_prefix(" ")

    99     inner_args = [for_leaf, fp, in_leaf, it]

   100     if test:

   101         test.set_prefix(" ")

   102         if_leaf = Leaf(token.NAME, "if")

   103         if_leaf.set_prefix(" ")

   104         inner_args.append(Node(syms.comp_if, [if_leaf, test]))

   105     inner = Node(syms.listmaker, [xp, Node(syms.comp_for, inner_args)])

   106     return Node(syms.atom,

   107                        [Leaf(token.LBRACE, "["),

   108                         inner,

   109                         Leaf(token.RBRACE, "]")])

   110 

   111 def FromImport(package_name, name_leafs):

   112     """ Return an import statement in the form:

   113         from package import name_leafs"""

   114     # XXX: May not handle dotted imports properly (eg, package_name='foo.bar')

   115     #assert package_name == '.' or '.' not in package_name, "FromImport has "\

   116     #       "not been tested with dotted package names -- use at your own "\

   117     #       "peril!"

   118 

   119     for leaf in name_leafs:

   120         # Pull the leaves out of their old tree

   121         leaf.remove()

   122 

   123     children = [Leaf(token.NAME, 'from'),

   124                 Leaf(token.NAME, package_name, prefix=" "),

   125                 Leaf(token.NAME, 'import', prefix=" "),

   126                 Node(syms.import_as_names, name_leafs)]

   127     imp = Node(syms.import_from, children)

   128     return imp

   129 

   130 

   131 ###########################################################

   132 ### Determine whether a node represents a given literal

   133 ###########################################################

   134 

   135 def is_tuple(node):

   136     """Does the node represent a tuple literal?"""

   137     if isinstance(node, Node) and node.children == [LParen(), RParen()]:

   138         return True

   139     return (isinstance(node, Node)

   140             and len(node.children) == 3

   141             and isinstance(node.children[0], Leaf)

   142             and isinstance(node.children[1], Node)

   143             and isinstance(node.children[2], Leaf)

   144             and node.children[0].value == "("

   145             and node.children[2].value == ")")

   146 

   147 def is_list(node):

   148     """Does the node represent a list literal?"""

   149     return (isinstance(node, Node)

   150             and len(node.children) > 1

   151             and isinstance(node.children[0], Leaf)

   152             and isinstance(node.children[-1], Leaf)

   153             and node.children[0].value == "["

   154             and node.children[-1].value == "]")

   155 

   156 

   157 ###########################################################

   158 ### Misc

   159 ###########################################################

   160 

   161 

   162 consuming_calls = set(["sorted", "list", "set", "any", "all", "tuple", "sum",

   163                        "min", "max"])

   164 

   165 def attr_chain(obj, attr):

   166     """Follow an attribute chain.

   167 

   168     If you have a chain of objects where a.foo -> b, b.foo-> c, etc,

   169     use this to iterate over all objects in the chain. Iteration is

   170     terminated by getattr(x, attr) is None.

   171 

   172     Args:

   173         obj: the starting object

   174         attr: the name of the chaining attribute

   175 

   176     Yields:

   177         Each successive object in the chain.

   178     """

   179     next = getattr(obj, attr)

   180     while next:

   181         yield next

   182         next = getattr(next, attr)

   183 

   184 p0 = """for_stmt< 'for' any 'in' node=any ':' any* >

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

   186      """

   187 p1 = """

   188 power<

   189     ( 'iter' | 'list' | 'tuple' | 'sorted' | 'set' | 'sum' |

   190       'any' | 'all' | (any* trailer< '.' 'join' >) )

   191     trailer< '(' node=any ')' >

   192     any*

   193 >

   194 """

   195 p2 = """

   196 power<

   197     'sorted'

   198     trailer< '(' arglist<node=any any*> ')' >

   199     any*

   200 >

   201 """

   202 pats_built = False

   203 def in_special_context(node):

   204     """ Returns true if node is in an environment where all that is required

   205         of it is being itterable (ie, it doesn't matter if it returns a list

   206         or an itterator).

   207         See test_map_nochange in test_fixers.py for some examples and tests.

   208         """

   209     global p0, p1, p2, pats_built

   210     if not pats_built:

   211         p1 = patcomp.compile_pattern(p1)

   212         p0 = patcomp.compile_pattern(p0)

   213         p2 = patcomp.compile_pattern(p2)

   214         pats_built = True

   215     patterns = [p0, p1, p2]

   216     for pattern, parent in zip(patterns, attr_chain(node, "parent")):

   217         results = {}

   218         if pattern.match(parent, results) and results["node"] is node:

   219             return True

   220     return False

   221 

   222 ###########################################################

   223 ### The following functions are to find bindings in a suite

   224 ###########################################################

   225 

   226 def make_suite(node):

   227     if node.type == syms.suite:

   228         return node

   229     node = node.clone()

   230     parent, node.parent = node.parent, None

   231     suite = Node(syms.suite, [node])

   232     suite.parent = parent

   233     return suite

   234 

   235 def does_tree_import(package, name, node):

   236     """ Returns true if name is imported from package at the

   237         top level of the tree which node belongs to.

   238         To cover the case of an import like 'import foo', use

   239         Null for the package and 'foo' for the name. """

   240     # Scamper up to the top level namespace

   241     while node.type != syms.file_input:

   242         assert node.parent, "Tree is insane! root found before "\

   243                            "file_input node was found."

   244         node = node.parent

   245 

   246     binding = find_binding(name, node, package)

   247     return bool(binding)

   248 

   249 _def_syms = set([syms.classdef, syms.funcdef])

   250 def find_binding(name, node, package=None):

   251     """ Returns the node which binds variable name, otherwise None.

   252         If optional argument package is supplied, only imports will

   253         be returned.

   254         See test cases for examples."""

   255     for child in node.children:

   256         ret = None

   257         if child.type == syms.for_stmt:

   258             if _find(name, child.children[1]):

   259                 return child

   260             n = find_binding(name, make_suite(child.children[-1]), package)

   261             if n: ret = n

   262         elif child.type in (syms.if_stmt, syms.while_stmt):

   263             n = find_binding(name, make_suite(child.children[-1]), package)

   264             if n: ret = n

   265         elif child.type == syms.try_stmt:

   266             n = find_binding(name, make_suite(child.children[2]), package)

   267             if n:

   268                 ret = n

   269             else:

   270                 for i, kid in enumerate(child.children[3:]):

   271                     if kid.type == token.COLON and kid.value == ":":

   272                         # i+3 is the colon, i+4 is the suite

   273                         n = find_binding(name, make_suite(child.children[i+4]), package)

   274                         if n: ret = n

   275         elif child.type in _def_syms and child.children[1].value == name:

   276             ret = child

   277         elif _is_import_binding(child, name, package):

   278             ret = child

   279         elif child.type == syms.simple_stmt:

   280             ret = find_binding(name, child, package)

   281         elif child.type == syms.expr_stmt:

   282             if _find(name, child.children[0]):

   283                 ret = child

   284 

   285         if ret:

   286             if not package:

   287                 return ret

   288             if ret.type in (syms.import_name, syms.import_from):

   289                 return ret

   290     return None

   291 

   292 _block_syms = set([syms.funcdef, syms.classdef, syms.trailer])

   293 def _find(name, node):

   294     nodes = [node]

   295     while nodes:

   296         node = nodes.pop()

   297         if node.type > 256 and node.type not in _block_syms:

   298             nodes.extend(node.children)

   299         elif node.type == token.NAME and node.value == name:

   300             return node

   301     return None

   302 

   303 def _is_import_binding(node, name, package=None):

   304     """ Will reuturn node if node will import name, or node

   305         will import * from package.  None is returned otherwise.

   306         See test cases for examples. """

   307 

   308     if node.type == syms.import_name and not package:

   309         imp = node.children[1]

   310         if imp.type == syms.dotted_as_names:

   311             for child in imp.children:

   312                 if child.type == syms.dotted_as_name:

   313                     if child.children[2].value == name:

   314                         return node

   315                 elif child.type == token.NAME and child.value == name:

   316                     return node

   317         elif imp.type == syms.dotted_as_name:

   318             last = imp.children[-1]

   319             if last.type == token.NAME and last.value == name:

   320                 return node

   321         elif imp.type == token.NAME and imp.value == name:

   322             return node

   323     elif node.type == syms.import_from:

   324         # unicode(...) is used to make life easier here, because

   325         # from a.b import parses to ['import', ['a', '.', 'b'], ...]

   326         if package and unicode(node.children[1]).strip() != package:

   327             return None

   328         n = node.children[3]

   329         if package and _find('as', n):

   330             # See test_from_import_as for explanation

   331             return None

   332         elif n.type == syms.import_as_names and _find(name, n):

   333             return node

   334         elif n.type == syms.import_as_name:

   335             child = n.children[2]

   336             if child.type == token.NAME and child.value == name:

   337                 return node

   338         elif n.type == token.NAME and n.value == name:

   339             return node

   340         elif package and n.type == token.STAR:

   341             return node

   342     return None