Python ast.iter_child_nodes() Examples

def walk(node, stop_at=tuple(), ignore=tuple()):
    """Walk through the children of an ast node.

    Args:
        node: an ast node
        stop_at: stop traversing through these nodes, including the matching
            node
        ignore: stop traversing through these nodes, excluding the matching
            node

    Returns: a generator of ast nodes
    """
    todo = deque([node])
    while todo:
        node = todo.popleft()
        if isinstance(node, ignore):
            # dequeue next node
            continue
        if not isinstance(node, stop_at):
            next_nodes = ast.iter_child_nodes(node)
            for n in next_nodes:
                todo.extend([n])
        yield node 

def get_config_comments(func):
    filename = inspect.getfile(func)
    func_body, line_offset = get_function_body(func)
    body_source = dedent_function_body(func_body)
    body_code = compile(body_source, filename, "exec", ast.PyCF_ONLY_AST)
    body_lines = body_source.split("\n")

    variables = {"seed": "the random seed for this experiment"}

    for ast_root in body_code.body:
        for ast_entry in [ast_root] + list(ast.iter_child_nodes(ast_root)):
            if isinstance(ast_entry, ast.Assign):
                # we found an assignment statement
                # go through all targets of the assignment
                # usually a single entry, but can be more for statements like:
                # a = b = 5
                for t in ast_entry.targets:
                    add_doc(t, variables, body_lines)

    return variables 

def _add_section(self, node):
        """
        Register the current node as a new context block
        """
        self._filldown(node.lineno)

        # push a new context onto stack
        self.context.append(node.name)
        self._update_current_context()

        for _ in map(self.visit, iter_child_nodes(node)):
            pass

        # restore current context
        self.context.pop()
        self._update_current_context() 

def __init__(self, *args, **kwargs):
        super(Source, self).__init__(*args, **kwargs)
        if self.tree and self.text:
            self.highlighted = ArgDefaultDict(
                lambda style: raw_highlight(self.text, style).splitlines()
            )
        else:
            self.lines = defaultdict(lambda: u'SOURCE IS UNAVAILABLE')
            self.highlighted = defaultdict(lambda: self.lines)
        self.statements = StatementsDict(self)
        self.nodes = []
        if self.tree:
            self.tree._depth = 0
            for node in ast.walk(self.tree):
                node._tree_index = len(self.nodes)
                self.nodes.append(node)
                for child in ast.iter_child_nodes(node):
                    child._depth = node._depth + 1 

def iter_children_ast(node):
  # Don't attempt to process children of JoinedStr nodes, which we can't fully handle yet.
  if is_joined_str(node):
    return

  if isinstance(node, ast.Dict):
    # override the iteration order: instead of <all keys>, <all values>,
    # yield keys and values in source order (key1, value1, key2, value2, ...)
    for (key, value) in zip(node.keys, node.values):
      if key is not None:
        yield key
      yield value
    return

  for child in ast.iter_child_nodes(node):
    # Skip singleton children; they don't reflect particular positions in the code and break the
    # assumptions about the tree consisting of distinct nodes. Note that collecting classes
    # beforehand and checking them in a set is faster than using isinstance each time.
    if child.__class__ not in SINGLETONS:
      yield child 

def locate_loop_body(module, loop):
    ends = set([ node.lineno 
                    for node in ast.walk(module) 
                    if hasattr(node, 'lineno') and node.lineno > loop.lineno ])

    starts = set()

    def visit(node):
        if hasattr(node, 'lineno') and node.lineno > loop.lineno:
            starts.add(node.lineno)
            if node.lineno in ends:
                ends.remove(node.lineno)
        for child in ast.iter_child_nodes(node):
            visit(child)

    for stmt in loop.body:
        visit(stmt)

    if len(ends) == 0: 
        return min(starts), -1

    return min(starts), min(ends) 

def check_for_b012(self, node):
        def _loop(node, bad_node_types):
            if isinstance(node, (ast.AsyncFunctionDef, ast.FunctionDef)):
                return

            if isinstance(node, (ast.While, ast.For)):
                bad_node_types = (ast.Return,)

            elif isinstance(node, bad_node_types):
                self.errors.append(B012(node.lineno, node.col_offset))

            for child in ast.iter_child_nodes(node):
                _loop(child, bad_node_types)

        for child in node.finalbody:
            _loop(child, (ast.Return, ast.Continue, ast.Break)) 

def diff(a, b):
        assert a is not None
        assert b is not None

        def _str_dist(i, j):
            return 0 if i == j else 1

        def _get_label(n):
            return type(n).__name__

        def _get_children(n):
            if not hasattr(n, 'children'):
                n.children = list(ast.iter_child_nodes(n))
            return n.children

        import zss
        res = zss.distance(a.func_node, b.func_node, _get_children,
                           lambda node: 0,  # insert cost
                           lambda node: _str_dist(_get_label(node), ''),  # remove cost
                           lambda _a, _b: _str_dist(_get_label(_a), _get_label(_b)), )  # update cost
        return res 

def parse_flags(synth_py_path="synth.py") -> typing.Dict[str, typing.Any]:
    """Extracts flags from a synth.py file.

    Keyword Arguments:
        synth_py_path {str or Path} -- Path to synth.py (default: {"synth.py"})

    Returns:
        typing.Dict[str, typing.Any] -- A map of all possible flags.  Flags not
          found in synth.py will be set to their default value.
    """
    flags = copy.copy(_SYNTH_PY_FLAGS)
    path = pathlib.Path(synth_py_path)
    try:
        module = ast.parse(path.read_text())
    except SyntaxError:
        return flags  # Later attempt to execute synth.py will give a clearer error message.

    for child in ast.iter_child_nodes(module):
        if isinstance(child, ast.Assign):
            for target in child.targets:
                if isinstance(target, ast.Name) and target.id in flags:
                    flags[target.id] = ast.literal_eval(child.value)
    return flags 

def get_version():
    module_path = os.path.join(
        os.path.dirname(__file__),
        'sqlalchemy_utc',
        'version.py')
    module_file = open(module_path)
    try:
        module_code = module_file.read()
    finally:
        module_file.close()
    tree = ast.parse(module_code, module_path)
    for node in ast.iter_child_nodes(tree):
        if not isinstance(node, ast.Assign) or len(node.targets) != 1:
            continue
        target, = node.targets
        if isinstance(target, ast.Name) and target.id == '__version__':
            return node.value.s 

def get_accessed(node):
    '''Get names, but ignore variables names to the left hand side
    That is to say, in case of
       a = b + 1
    we consider b as "being accessed", while a is not.
    '''
    if isinstance(node, ast.Assign):
        return get_accessed(node.value)
    elif isinstance(node, ast.Name):
        return {node.id}
    names = set()
    if isinstance(node, list):
        for x in node:
            names |= get_accessed(x)
    else:
        for x in ast.iter_child_nodes(node):
            names |= get_accessed(x)
    return names 

def _process_child_nodes(
    node: ast.AST,
    increment_by: int,
    complexity_calculator: Callable[[ast.AST, int], int],
) -> int:
    child_complexity = 0

    for node_num, child_node in enumerate(ast.iter_child_nodes(node)):
        if isinstance(node, ast.Try):
            if node_num == 1:
                increment_by += 1  # add +1 for all try nodes except body
            if node_num:
                child_complexity += max(1, increment_by)
        child_complexity += complexity_calculator(
            child_node,
            increment_by,
        )

    return child_complexity 

def _set_parent(tree: ast.AST) -> ast.AST:
    """
    Sets parents for all nodes that do not have this prop.

    This step is required due to how `flake8` works.
    It does not set the same properties as `ast` module.

    This function was the cause of `issue-112`. Twice.
    Since the ``0.6.1`` we use ``'wps_parent'`` with a prefix.
    This should fix the issue with conflicting plugins.

    .. versionchanged:: 0.0.11
    .. versionchanged:: 0.6.1

    """
    for statement in ast.walk(tree):
        for child in ast.iter_child_nodes(statement):
            setattr(child, 'wps_parent', statement)  # noqa: B010
    return tree 

def _apply_if_statement(statement: ast.If) -> None:
    """We need to add extra properties to ``if`` conditions."""
    for child in ast.iter_child_nodes(statement):
        if isinstance(child, ast.If):
            if child in statement.orelse:
                setattr(statement, 'wps_if_chained', True)  # noqa: B010
                setattr(child, 'wps_if_chain', statement)  # noqa: B010 

def _chained_item(iterator: ast.AST) -> Optional[ast.Call]:
    children = list(ast.iter_child_nodes(iterator))
    if isinstance(children[0], ast.Call):
        return children[0]
    return None 

def generic_visit(self, node):
        if hasattr(node, 'lineno'):
            self._filldown(node.lineno + 1)

        for _ in map(self.visit, iter_child_nodes(node)):
            pass 

def _find_closest_containing_node(tree, text_range):
    # first look among children
    for child in ast.iter_child_nodes(tree):
        result = _find_closest_containing_node(child, text_range)
        if result is not None:
            return result

    # no suitable child was found
    if hasattr(tree, "lineno") and range_contains_smaller(
        TextRange(tree.lineno, tree.col_offset, tree.end_lineno, tree.end_col_offset), text_range
    ):
        return tree
    # nope
    else:
        return None 

def _list_tasks():
    """
    Fetches a list of all valid tasks that may be run, and the args they
    accept. Does not actually import the task module to prevent errors if a
    user does not have the dependencies installed for every task.

    :return:
        A list of 2-element tuples:
         0: a unicode string of the task name
         1: a list of dicts containing the parameter definitions
    """

    out = []
    dev_path = os.path.join(package_root, 'dev')
    for fname in sorted(os.listdir(dev_path)):
        if fname.startswith('.') or fname.startswith('_'):
            continue
        if not fname.endswith('.py'):
            continue
        name = fname[:-3]
        args = ()

        full_path = os.path.join(package_root, 'dev', fname)
        with open(full_path, 'rb') as f:
            full_code = f.read()
            if sys.version_info >= (3,):
                full_code = full_code.decode('utf-8')

        task_node = ast.parse(full_code, filename=full_path)
        for node in ast.iter_child_nodes(task_node):
            if isinstance(node, _ast.Assign):
                if len(node.targets) == 1 \
                        and isinstance(node.targets[0], _ast.Name) \
                        and node.targets[0].id == 'run_args':
                    args = ast.literal_eval(node.value)
                    break

        out.append((name, args))
    return out 

def default(self, node, *args):
        for child in iter_child_nodes(node):
            self.dispatch(child, *args) 

def generic_visit(self, node):
        # visit all children
        super(ast.NodeVisitor, self).generic_visit(node)
        # bubble up code_nodes from the children if any (cleaning up children
        # at the same time to keep sound asts)
        node.code_nodes = []
        for child in ast.iter_child_nodes(node):
            node.code_nodes = node.code_nodes + child.code_nodes
            del(child.code_nodes) 

def test_iter_child_nodes(self):
        node = ast.parse("spam(23, 42, eggs='leek')", mode='eval')
        self.assertEqual(len(list(ast.iter_child_nodes(node.body))), 4)
        iterator = ast.iter_child_nodes(node.body)
        self.assertEqual(iterator.next().id, 'spam')
        self.assertEqual(iterator.next().n, 23)
        self.assertEqual(iterator.next().n, 42)
        self.assertEqual(ast.dump(iterator.next()),
            "keyword(arg='eggs', value=Str(s='leek'))"
        ) 

def test_iter_child_nodes(self):
        node = ast.parse("spam(23, 42, eggs='leek')", mode='eval')
        self.assertEqual(len(list(ast.iter_child_nodes(node.body))), 4)
        iterator = ast.iter_child_nodes(node.body)
        self.assertEqual(next(iterator).id, 'spam')
        self.assertEqual(next(iterator).n, 23)
        self.assertEqual(next(iterator).n, 42)
        self.assertEqual(ast.dump(next(iterator)),
            "keyword(arg='eggs', value=Str(s='leek'))"
        ) 

def test_iter_child_nodes(self):
        node = ast.parse("spam(23, 42, eggs='leek')", mode='eval')
        self.assertEqual(len(list(ast.iter_child_nodes(node.body))), 4)
        iterator = ast.iter_child_nodes(node.body)
        self.assertEqual(next(iterator).id, 'spam')
        self.assertEqual(next(iterator).n, 23)
        self.assertEqual(next(iterator).n, 42)
        self.assertEqual(ast.dump(next(iterator)),
            "keyword(arg='eggs', value=Str(s='leek'))"
        ) 

def fix_missing_locations(node):
    """
    When you compile a node tree with compile(), the compiler expects lineno
    and col_offset attributes for every node that supports them.  This is
    rather tedious to fill in for generated nodes, so this helper adds these
    attributes recursively where not already set, by setting them to the values
    of the parent node.  It works recursively starting at *node*.
    """
    def _fix(node, lineno, col_offset, end_lineno, end_col_offset):
        if 'lineno' in node._attributes:
            if not hasattr(node, 'lineno'):
                node.lineno = lineno
            else:
                lineno = node.lineno
        if 'end_lineno' in node._attributes:
            if not hasattr(node, 'end_lineno'):
                node.end_lineno = end_lineno
            else:
                end_lineno = node.end_lineno
        if 'col_offset' in node._attributes:
            if not hasattr(node, 'col_offset'):
                node.col_offset = col_offset
            else:
                col_offset = node.col_offset
        if 'end_col_offset' in node._attributes:
            if not hasattr(node, 'end_col_offset'):
                node.end_col_offset = end_col_offset
            else:
                end_col_offset = node.end_col_offset
        for child in iter_child_nodes(node):
            _fix(child, lineno, col_offset, end_lineno, end_col_offset)
    _fix(node, 1, 0, 1, 0)
    return node 

def walk_function_body(self, node):
        def _loop(parent, node):
            if isinstance(node, (ast.AsyncFunctionDef, ast.FunctionDef)):
                return
            yield parent, node
            for child in ast.iter_child_nodes(node):
                yield from _loop(node, child)

        for child in node.body:
            yield from _loop(node, child) 

def set_location(node, lineno, col_offset):
    """Set node location information recursively."""

    def _fix(node, lineno, col_offset):
        if "lineno" in node._attributes:
            node.lineno = lineno
        if "col_offset" in node._attributes:
            node.col_offset = col_offset
        for child in ast.iter_child_nodes(node):
            _fix(child, lineno, col_offset)

    _fix(node, lineno, col_offset)
    return node 

def is_interesting_expression(node):
    # type: (ast.AST) -> bool
    """
    If this expression has a value that may not be exactly what it looks like,
    return True. Put differently, return False if this is just a literal.
    """
    return (isinstance(node, ast.expr) and
            not (isinstance(node, (ast.Num, ast.Str, getattr(ast, 'NameConstant', ()))) or
                 isinstance(getattr(node, 'ctx', None),
                            (ast.Store, ast.Del)) or
                 (isinstance(node, ast.UnaryOp) and
                  isinstance(node.op, (ast.UAdd, ast.USub)) and
                  isinstance(node.operand, ast.Num)) or
                 (isinstance(node, (ast.List, ast.Tuple, ast.Dict)) and
                  not any(is_interesting_expression(n) for n in ast.iter_child_nodes(node))))) 

def set_basic_node_attributes(self):
        self.nodes = []  # type: List[ast.AST]
        for node in ast.walk(self.root):  # type: ast.AST
            for child in ast.iter_child_nodes(node):
                child.parent = node
            node._tree_index = len(self.nodes)
            self.nodes.append(node)

        # Mark __future__ imports and anything before (i.e. module docstrings)
        # to be ignored by the AST transformer
        for i, stmt in enumerate(self.root.body):
            if is_future_import(stmt):
                for s in self.root.body[:i + 1]:
                    for node in ast.walk(s):
                        node._visit_ignore = True 

def visit_FunctionDef(self, node, name=None):
        name = name or node.name
        self.stack.append(name)
        self.qualnames.setdefault((name, node.lineno), ".".join(self.stack))

        self.stack.append('<locals>')
        if isinstance(node, ast.Lambda):
            children = [node.body]
        else:
            children = node.body
        for child in children:
            self.visit(child)
        self.stack.pop()
        self.stack.pop()

        # Find lambdas in the function definition outside the body,
        # e.g. decorators or default arguments
        # Based on iter_child_nodes
        for field, child in ast.iter_fields(node):
            if field == 'body':
                continue
            if isinstance(child, ast.AST):
                self.visit(child)
            elif isinstance(child, list):
                for grandchild in child:
                    if isinstance(grandchild, ast.AST):
                        self.visit(grandchild) 

def get_testcases(paths):
    """Walk each path in ``paths`` and return the test cases found.

    :param path: List o directories to find test modules and test cases.
    :return: A dict mapping a test module path and its test cases.
    """
    testmodules = []
    for path in paths:
        if os.path.isfile(path):
            if is_test_module(os.path.basename(path)):
                testmodules.append(path)
            continue
        for dirpath, _, filenames in os.walk(path):
            for filename in filenames:
                if is_test_module(filename):
                    testmodules.append(os.path.join(dirpath, filename))
    testcases = collections.OrderedDict()
    for testmodule in testmodules:
        testcases[testmodule] = []
        with open(testmodule) as handler:
            root = ast.parse(handler.read())
            root.path = testmodule
            for node in ast.iter_child_nodes(root):
                if isinstance(node, ast.ClassDef):
                    testcases[testmodule].extend([
                        TestFunction(subnode, node, root)
                        for subnode in ast.iter_child_nodes(node)
                        if isinstance(subnode, ast.FunctionDef) and
                        subnode.name.startswith('test_')
                    ])
                elif (isinstance(node, ast.FunctionDef) and
                      node.name.startswith('test_')):
                    # Module's test functions
                    testcases[testmodule].append(
                        TestFunction(node, testmodule=root)
                    )
    return testcases