Python ast.iter_fields() Examples

def ast_to_string(ast_node, line_indent=''):
    next_line_indent = line_indent + '  '

    if isinstance(ast_node, ast.AST):
        return (ast_node.__class__.__name__
                + '('
                + ','.join('\n' + next_line_indent + field_name + ' = ' + ast_to_string(child_node, next_line_indent)
                           for field_name, child_node in ast.iter_fields(ast_node))
                + ')')
    elif isinstance(ast_node, list):
        return ('['
                + ','.join('\n' + next_line_indent + ast_to_string(child_node, next_line_indent)
                           for child_node in ast_node)
                + ']')
    else:
        return repr(ast_node) 

def test_iter_fields(self):
        node = ast.parse('foo()', mode='eval')
        d = dict(ast.iter_fields(node.body))
        self.assertEqual(d.pop('func').id, 'foo')

        #XXX: tests for equality between astlist and regular lists not
        #     working, breaking this test up into its components.
        #self.assertEqual(d, {'keywords': [], 'kwargs': None,
        #                     'args': [], 'starargs': None})
        assert len(d) == 4
        assert d['keywords'] is not None
        assert len(d['keywords']) == 0
        assert d['args'] is not None
        assert len(d['args']) == 0
        assert d['kwargs'] is None
        assert d['starargs'] is None 

def parse(self, path):
        schemas = None
        file = open(path, 'r')
        data = file.read()
        root = ast.parse(data)

        try:
            for node in ast.walk(root):
                if isinstance(node, ast.Assign):
                    if hasattr(node.targets[0], 'id') and node.targets[0].id in schema_names:
                        schema_name = node.targets[0].id
                        for name, val in ast.iter_fields(node):
                            if isinstance(val, ast.Dict):
                                if not schemas:
                                    schemas = {}
                                schemas[schema_name] = self.convert(val)
        except:
            schemas = None
        file.close()
        return schemas 

def visit_Assign(self, node):
        for key, val in ast.iter_fields(node):
            if key == 'targets':
                for subnode in val:
                    if type(subnode) is ast.Tuple:
                        for subsubnode in subnode.elts:
                            crawler = myextract()
                            crawler.visit(subsubnode)
                            self.targets[fix_assign.sub(r'\1', ".".join(reversed(crawler.stack)))] = 1
                    else:
                        crawler = myextract()
                        crawler.visit(subnode)
                        self.targets[fix_assign.sub(r'\1', ".".join(reversed(crawler.stack)))] = 1
        self.depth += 1
        #ast.NodeVisitor.generic_visit(self, node)
        self.generic_visit(node)
        self.depth -= 1 

def generic_visit(self, node: ast.AST) -> None:
        """Called if no explicit visitor function exists for a node."""
        for _field, value in ast.iter_fields(node):
            if self.should_type_check:
                break

            if isinstance(value, list):
                for item in value:
                    if self.should_type_check:
                        break
                    if isinstance(item, ast.AST):
                        self.visit(item)
            elif isinstance(value, ast.AST):
                self.visit(value)


# Generic mypy error 

def dumpTree (self):
        utils.log (False, 'Dumping syntax tree for module: {}\n', self.sourcePath)

        def walk (name, value, tabLevel):
            self.treeFragments .append ('\n{0}{1}: {2} '.format (tabLevel * '\t', name, type (value).__name__ ))
            if isinstance (value, ast.AST):
                for field in ast.iter_fields (value):
                    walk (field [0], field [1], tabLevel + 1)
            elif isinstance (value, list):
                for element in value:
                    walk ('element', element, tabLevel + 1)
            else:
                self.treeFragments.append ('= {0}'.format (value))

        self.treeFragments = []
        walk ('file', self.parseTree, 0)
        self.textTree = ''.join (self.treeFragments) [1:]

        with utils.create (self.treePath) as treeFile:
            treeFile.write (self.textTree) 

def test_iter_fields(self):
        node = ast.parse('foo()', mode='eval')
        d = dict(ast.iter_fields(node.body))
        self.assertEqual(d.pop('func').id, 'foo')

        #XXX: tests for equality between astlist and regular lists not
        #     working, breaking this test up into its components.
        #self.assertEqual(d, {'keywords': [], 'kwargs': None,
        #                     'args': [], 'starargs': None})
        assert len(d) == 4
        assert d['keywords'] is not None
        assert len(d['keywords']) == 0
        assert d['args'] is not None
        assert len(d['args']) == 0
        assert d['kwargs'] is None
        assert d['starargs'] is None 

def test_iter_fields(self):
        node = ast.parse('foo()', mode='eval')
        d = dict(ast.iter_fields(node.body))
        self.assertEqual(d.pop('func').id, 'foo')

        #XXX: tests for equality between astlist and regular lists not
        #     working, breaking this test up into its components.
        #self.assertEqual(d, {'keywords': [], 'kwargs': None,
        #                     'args': [], 'starargs': None})
        assert len(d) == 4
        assert d['keywords'] is not None
        assert len(d['keywords']) == 0
        assert d['args'] is not None
        assert len(d['args']) == 0
        assert d['kwargs'] is None
        assert d['starargs'] is None 

def compare(computed, expected):
    """
    Assert that two AST nodes are the same.
    """
    assert type(computed) == type(expected)

    if isinstance(computed, list):
        for cv, ev in zip_longest(computed, expected):
            compare(cv, ev)
        return

    if not isinstance(computed, AST):
        assert computed == expected
        return

    for (cn, cv), (en, ev) in zip_longest(*map(iter_fields,
                                               (computed, expected))):
        assert cn == en
        compare(cv, ev) 

def generic_visit(self, node):
        """Drive the visitor."""
        for _, value in ast.iter_fields(node):
            if isinstance(value, list):
                max_idx = len(value) - 1
                for idx, item in enumerate(value):
                    if isinstance(item, ast.AST):
                        if idx < max_idx:
                            item._bandit_sibling = value[idx + 1]
                        else:
                            item._bandit_sibling = None
                        item._bandit_parent = node

                        if self.pre_visit(item):
                            self.visit(item)
                            self.generic_visit(item)
                            self.post_visit(item)

            elif isinstance(value, ast.AST):
                value._bandit_sibling = None
                value._bandit_parent = node
                if self.pre_visit(value):
                    self.visit(value)
                    self.generic_visit(value)
                    self.post_visit(value) 

def fields_same(n1: ast.AST, n2: ast.AST) -> bool:
    for (a1, v1), (a2, v2) in zip(ast.iter_fields(n1), ast.iter_fields(n2)):
        # ignore ast attributes, they'll be covered by walk
        if a1 != a2:
            return False
        elif _all_isinstance((v1, v2), ast.AST):
            continue
        elif _all_isinstance((v1, v2), (list, tuple)):
            if len(v1) != len(v2):
                return False
            # ignore sequences which are all-ast, they'll be covered by walk
            elif _all_isinstance(v1, ast.AST) and _all_isinstance(v2, ast.AST):
                continue
            elif v1 != v2:
                return False
        elif v1 != v2:
            return False
    return True 

def test_iter_fields(self):
        node = ast.parse('foo()', mode='eval')
        d = dict(ast.iter_fields(node.body))
        self.assertEqual(d.pop('func').id, 'foo')
        self.assertEqual(d, {'keywords': [], 'kwargs': None,
                             'args': [], 'starargs': None}) 

def visit_AST(self, node):
        for key, value in ast.iter_fields(node):
            if value is None:
                continue
            elif isinstance(value, list):
                list_id = self._id(key)
                self.graph[self._id(node)].append(list_id)
                for item in value:
                    self.graph[list_id].append(self._id(item))
                    self.visit(item)
            else:
                self.graph[self._id(node)].append(self._id(value))
                self.visit(value) 

def test_iter_fields(self):
        node = ast.parse('foo()', mode='eval')
        d = dict(ast.iter_fields(node.body))
        self.assertEqual(d.pop('func').id, 'foo')
        self.assertEqual(d, {'keywords': [], 'kwargs': None,
                             'args': [], 'starargs': None}) 

def test_iter_fields(self):
        node = ast.parse('foo()', mode='eval')
        d = dict(ast.iter_fields(node.body))
        self.assertEqual(d.pop('func').id, 'foo')
        self.assertEqual(d, {'keywords': [], 'args': []}) 

def test_iter_fields(self):
        node = ast.parse('foo()', mode='eval')
        d = dict(ast.iter_fields(node.body))
        self.assertEqual(d.pop('func').id, 'foo')
        self.assertEqual(d, {'keywords': [], 'kwargs': None,
                             'args': [], 'starargs': None}) 

def print_ast(node, level):
    print(" " * level, node)
    for name, child in ast.iter_fields(node):
        if isinstance(child, ast.AST):
            print(" " * level, name, ":")
            print_ast(child, level + 1)
        elif isinstance(child, list):
            print(" " * level, name, ":[")
            for elem in child:
                print_ast(elem, level + 1)
            print(" " * level, "]")
        else:
            pass
            # print(" " * level, "OOOO", node) 

def _insert_statement_markers(self, root):
        # find lists of statements and insert before/after markers for each statement
        for name, value in ast.iter_fields(root):
            if isinstance(root, ast.Try) and name == "handlers":
                # contains statements but is not statement itself
                for handler in value:
                    self._insert_statement_markers(handler)
            elif isinstance(value, ast.AST):
                self._insert_statement_markers(value)
            elif isinstance(value, list):
                if len(value) > 0:
                    new_list = []
                    for node in value:
                        if self._should_instrument_as_statement(node):
                            # self._debug("EBFOMA", node)
                            # add before marker
                            new_list.append(
                                self._create_statement_marker(node, BEFORE_STATEMENT_MARKER)
                            )

                        # original statement
                        if self._should_instrument_as_statement(node):
                            self._insert_statement_markers(node)
                        new_list.append(node)

                        if (
                            self._should_instrument_as_statement(node)
                            and "skipexport" not in node.tags
                        ):
                            # add after marker
                            new_list.append(
                                self._create_statement_marker(node, AFTER_STATEMENT_MARKER)
                            )
                    setattr(root, name, new_list) 

def generic_visit(self, node):
        """Called if no explicit visitor function exists for a node."""
        for _, value in ast.iter_fields(node):
            if isinstance(value, list):
                self._handle_ast_list(value)
                for item in value:
                    if isinstance(item, ast.AST):
                        self.visit(item)
            elif isinstance(value, ast.AST):
                self.visit(value) 

def ast_visit(node, level=0):
	print('  ' * level + str_node(node))
	for field, value in ast.iter_fields(node):
		if isinstance(value, list):
			for item in value:
				if isinstance(item, ast.AST):
					ast_visit(item, level=level+1)
		elif isinstance(value, ast.AST):
			ast_visit(value, level=level+1) 

def str_node(node):
	if isinstance(node, ast.AST):
		fields = [(name, str_node(val)) for name, val in ast.iter_fields(node) if name not in ('left', 'right')]
		rv = '%s(%s' % (node.__class__.__name__, ', '.join('%s=%s' % field for field in fields))
		return rv + ')'
	else:
		return repr(node) 

def clone_node(node: AST_T, **updated_args) -> AST_T:
    """Returns a shallow copy of an AST node with the specified attributes updated."""
    args = dict(ast.iter_fields(node))
    args.update(updated_args)
    ast_class = type(node)
    return ast_class(**args) 

def remove_child(parent, child):
  for _, field_value in ast.iter_fields(parent):
    if isinstance(field_value, list) and child in field_value:
      field_value.remove(child)
      return
  raise errors.InvalidAstError('Unable to find list containing child %r on '
                               'parent node %r' % (child, parent)) 

def test_iter_fields(self):
        node = ast.parse('foo()', mode='eval')
        d = dict(ast.iter_fields(node.body))
        self.assertEqual(d.pop('func').id, 'foo')
        self.assertEqual(d, {'keywords': [], 'kwargs': None,
                             'args': [], 'starargs': None}) 

def visit_AugAssign(self, node):
        for key, val in ast.iter_fields(node):
            if key == 'target':
                crawler = myextract()
                crawler.visit(val)
                self.targets[fix_assign.sub(r'\1', ".".join(reversed(crawler.stack)))] = 1
        self.depth += 1
        #ast.NodeVisitor.generic_visit(self, node)
        self.generic_visit(node)
        self.depth -= 1 

def visit_AnnAssign(self, node):
        for key, val in ast.iter_fields(node):
            if key == 'target':
                crawler = myextract()
                crawler.visit(val)
                self.targets[fix_assign.sub(r'\1', ".".join(reversed(crawler.stack)))] = 1
        self.depth += 1
        #ast.NodeVisitor.generic_visit(self, node)
        self.generic_visit(node)
        self.depth -= 1 

def test_iter_fields(self):
        node = ast.parse('foo()', mode='eval')
        d = dict(ast.iter_fields(node.body))
        self.assertEqual(d.pop('func').id, 'foo')
        self.assertEqual(d, {'keywords': [], 'kwargs': None,
                             'args': [], 'starargs': None}) 

def extract_option_from_arg_list(options: str, optname: str,
                                 default_value: None) -> Tuple[Any, str]:
    if not options:
        return default_value, options
    try:
        args = list(ast.iter_fields(ast.parse(f"f({options})",
                                              mode='eval')))[0][1].keywords
        for field in args:
            if field.arg == optname:
                try:
                    value = eval(
                        compile(
                            ast.Expression(body=field.value),
                            filename="<ast>",
                            mode="eval"))
                    new_options = ','.join([
                        x for x in options.split(',')
                        if not x.strip().startswith(optname)
                    ])
                    return value, new_options.strip()
                except Exception:
                    raise ValueError(
                        f"A constant value is expected for option {optname}: {options} provided."
                    )
        return default_value, options
    except SyntaxError:
        raise ValueError(
            f"Expect a list of keyword arguments: {options} provided") 

def dump(node, annotate_fields=True, include_attributes=False, indent='  '):
    """Return a formatted dump of the tree in *node*.

    This is mainly useful for debugging purposes.  The returned string will
    show the names and the values for fields.  This makes the code impossible
    to evaluate, so if evaluation is wanted *annotate_fields* must be set to
    False.  Attributes such as line numbers and column offsets are not dumped
    by default.  If this is wanted, *include_attributes* can be set to True.
    """
    def _format(node, level=0):
        if isinstance(node, AST):
            fields = [(a, _format(b, level + 1)) for a, b in iter_fields(node)]
            if include_attributes and node._attributes:
                fields.extend([(a, _format(getattr(node, a), level + 1))
                               for a in node._attributes])
            return ''.join([
                node.__class__.__name__,
                '(\n' + indent + indent * level if fields else '(',
                (',\n' + indent + indent * level).join(('%s=%s' % field for field in fields)
                          if annotate_fields else
                          (b for a, b in fields)),
                ')'])
        elif isinstance(node, list):
            lines = ['[']
            lines.extend((indent * (level + 2) + _format(x, level + 2) + ','
                         for x in node))
            if len(lines) > 1:
                lines.append(indent * (level + 1) + ']')
            else:
                lines[-1] += ']'
            return '\n'.join(lines)
        return repr(node)

    if not isinstance(node, AST):
        raise TypeError('expected AST, got %r' % node.__class__.__name__)
    return _format(node) 

def get_setter(node):
    parent = node.parent
    for name, field in ast.iter_fields(parent):
        if field is node:
            return lambda new_node: setattr(parent, name, new_node)
        elif isinstance(field, list):
            for i, item in enumerate(field):
                if item is node:
                    def setter(new_node):
                        field[i] = new_node

                    return setter