Adding new stuff

vim-plugins/python-mode/pymode/libs/astroid/test_utils.py

@@ -0,0 +1,218 @@
+"""Utility functions for test code that uses astroid ASTs as input."""
+import functools
+import sys
+import textwrap
+
+from astroid import nodes
+from astroid import builder
+# The name of the transient function that is used to
+# wrap expressions to be extracted when calling
+# extract_node.
+_TRANSIENT_FUNCTION = '__'
+
+# The comment used to select a statement to be extracted
+# when calling extract_node.
+_STATEMENT_SELECTOR = '#@'
+
+
+def _extract_expressions(node):
+    """Find expressions in a call to _TRANSIENT_FUNCTION and extract them.
+
+    The function walks the AST recursively to search for expressions that
+    are wrapped into a call to _TRANSIENT_FUNCTION. If it finds such an
+    expression, it completely removes the function call node from the tree,
+    replacing it by the wrapped expression inside the parent.
+
+    :param node: An astroid node.
+    :type node:  astroid.bases.NodeNG
+    :yields: The sequence of wrapped expressions on the modified tree
+    expression can be found.
+    """
+    if (isinstance(node, nodes.CallFunc)
+            and isinstance(node.func, nodes.Name)
+            and node.func.name == _TRANSIENT_FUNCTION):
+        real_expr = node.args[0]
+        real_expr.parent = node.parent
+        # Search for node in all _astng_fields (the fields checked when
+        # get_children is called) of its parent. Some of those fields may
+        # be lists or tuples, in which case the elements need to be checked.
+        # When we find it, replace it by real_expr, so that the AST looks
+        # like no call to _TRANSIENT_FUNCTION ever took place.
+        for name in node.parent._astroid_fields:
+            child = getattr(node.parent, name)
+            if isinstance(child, (list, tuple)):
+                for idx, compound_child in enumerate(child):
+                    if compound_child is node:
+                        child[idx] = real_expr
+            elif child is node:
+                setattr(node.parent, name, real_expr)
+        yield real_expr
+    else:
+        for child in node.get_children():
+            for result in _extract_expressions(child):
+                yield result
+
+
+def _find_statement_by_line(node, line):
+    """Extracts the statement on a specific line from an AST.
+
+    If the line number of node matches line, it will be returned;
+    otherwise its children are iterated and the function is called
+    recursively.
+
+    :param node: An astroid node.
+    :type node: astroid.bases.NodeNG
+    :param line: The line number of the statement to extract.
+    :type line: int
+    :returns: The statement on the line, or None if no statement for the line
+      can be found.
+    :rtype:  astroid.bases.NodeNG or None
+    """
+    if isinstance(node, (nodes.Class, nodes.Function)):
+        # This is an inaccuracy in the AST: the nodes that can be
+        # decorated do not carry explicit information on which line
+        # the actual definition (class/def), but .fromline seems to
+        # be close enough.
+        node_line = node.fromlineno
+    else:
+        node_line = node.lineno
+
+    if node_line == line:
+        return node
+
+    for child in node.get_children():
+        result = _find_statement_by_line(child, line)
+        if result:
+            return result
+
+    return None
+
+def extract_node(code, module_name=''):
+    """Parses some Python code as a module and extracts a designated AST node.
+
+    Statements:
+     To extract one or more statement nodes, append #@ to the end of the line
+
+     Examples:
+       >>> def x():
+       >>>   def y():
+       >>>     return 1 #@
+
+       The return statement will be extracted.
+
+       >>> class X(object):
+       >>>   def meth(self): #@
+       >>>     pass
+
+      The funcion object 'meth' will be extracted.
+
+    Expressions:
+     To extract arbitrary expressions, surround them with the fake
+     function call __(...). After parsing, the surrounded expression
+     will be returned and the whole AST (accessible via the returned
+     node's parent attribute) will look like the function call was
+     never there in the first place.
+
+     Examples:
+       >>> a = __(1)
+
+       The const node will be extracted.
+
+       >>> def x(d=__(foo.bar)): pass
+
+       The node containing the default argument will be extracted.
+
+       >>> def foo(a, b):
+       >>>   return 0 < __(len(a)) < b
+
+       The node containing the function call 'len' will be extracted.
+
+    If no statements or expressions are selected, the last toplevel
+    statement will be returned.
+
+    If the selected statement is a discard statement, (i.e. an expression
+    turned into a statement), the wrapped expression is returned instead.
+
+    For convenience, singleton lists are unpacked.
+
+    :param str code: A piece of Python code that is parsed as
+    a module. Will be passed through textwrap.dedent first.
+    :param str module_name: The name of the module.
+    :returns: The designated node from the parse tree, or a list of nodes.
+    :rtype: astroid.bases.NodeNG, or a list of nodes.
+    """
+    def _extract(node):
+        if isinstance(node, nodes.Discard):
+            return node.value
+        else:
+            return node
+
+    requested_lines = []
+    for idx, line in enumerate(code.splitlines()):
+        if line.strip().endswith(_STATEMENT_SELECTOR):
+            requested_lines.append(idx + 1)
+
+    tree = build_module(code, module_name=module_name)
+    extracted = []
+    if requested_lines:
+        for line in requested_lines:
+            extracted.append(_find_statement_by_line(tree, line))
+
+    # Modifies the tree.
+    extracted.extend(_extract_expressions(tree))
+
+    if not extracted:
+        extracted.append(tree.body[-1])
+
+    extracted = [_extract(node) for node in extracted]
+    if len(extracted) == 1:
+        return extracted[0]
+    else:
+        return extracted
+
+
+def build_module(code, module_name='', path=None):
+    """Parses a string module with a builder.
+    :param code: The code for the module.
+    :type code: str
+    :param module_name: The name for the module
+    :type module_name: str
+    :param path: The path for the module
+    :type module_name: str
+    :returns: The module AST.
+    :rtype:  astroid.bases.NodeNG
+    """
+    code = textwrap.dedent(code)
+    return builder.AstroidBuilder(None).string_build(code, modname=module_name, path=path)
+
+
+def require_version(minver=None, maxver=None):
+    """ Compare version of python interpreter to the given one. Skip the test
+    if older.
+    """
+    def parse(string, default=None):
+        string = string or default
+        try:
+            return tuple(int(v) for v in string.split('.'))
+        except ValueError:
+            raise ValueError('%s is not a correct version : should be X.Y[.Z].' % version)
+
+    def check_require_version(f):
+        current = sys.version_info[:3]
+        if parse(minver, "0") < current <= parse(maxver, "4"):
+            return f
+        else:
+            str_version = '.'.join(str(v) for v in sys.version_info)
+            @functools.wraps(f)
+            def new_f(self, *args, **kwargs):
+                if minver is not None:
+                    self.skipTest('Needs Python > %s. Current version is %s.' % (minver, str_version))
+                elif maxver is not None:
+                    self.skipTest('Needs Python <= %s. Current version is %s.' % (maxver, str_version))
+            return new_f
+
+
+    return check_require_version
+
+def get_name_node(start_from, name, index=0):
+    return [n for n in start_from.nodes_of_class(nodes.Name) if n.name == name][index]