Python ast.walk() Examples

def compile_expression(exp):
    cached = cache.get(exp)
    if cached is not None:
        return cached
    _exp = ast.parse(exp)
    nodes = [node for node in ast.walk(_exp)]
    if len(nodes) < 2 or not isinstance(nodes[1], ast.Expr):
        raise ExpressionError("%s is not Expression" % exp)
    for node in nodes:
        if isinstance(node, ast.Call):
            raise ExpressionError("Call method is forbidden")
        if isinstance(node, ast.Lambda):
            raise ExpressionError("Lambda is strongly forbidden")
    result = compile(exp, '<string>', mode='eval')
    cache[exp] = result
    return result 

def walk_tree(node):
    try:
        d = node.__dict__
    except AttributeError:
        d = {}
        yield node
    for key, value in d.items():
        if isinstance(value, list):
            for val in value:
                for _ in ast.walk(val):
                    yield _
        elif issubclass(type(value), ast.AST):
            for _ in walk_tree(value):
                yield _
        else:
            yield value 

def get_source_file_contents(self, extension_list, directories_to_skip=None):
        file_paths = []
        file_contents = []
        directories_to_skip = directories_to_skip or []
        for dirname, directories_list, filenames in os.walk(self.path, topdown=True):
            directories_list[:] = [
                d for d in directories_list
                if d not in directories_to_skip
            ]
            for filename in filenames:
                extension = os.path.splitext(filename)[1]
                if extension in extension_list:
                    file_paths.append(os.path.join(dirname, filename))
        for file_path in file_paths:
            with open(file_path, 'r', encoding='utf-8') as file_handler:
                file_contents.append(file_handler.read())
        source_file_contents = list(zip(file_paths, file_contents))
        return source_file_contents 

def get_statement_startend2(lineno, node):
    import ast
    # flatten all statements and except handlers into one lineno-list
    # AST's line numbers start indexing at 1
    l = []
    for x in ast.walk(node):
        if isinstance(x, _ast.stmt) or isinstance(x, _ast.ExceptHandler):
            l.append(x.lineno - 1)
            for name in "finalbody", "orelse":
                val = getattr(x, name, None)
                if val:
                    # treat the finally/orelse part as its own statement
                    l.append(val[0].lineno - 1 - 1)
    l.sort()
    insert_index = bisect_right(l, lineno)
    start = l[insert_index - 1]
    if insert_index >= len(l):
        end = None
    else:
        end = l[insert_index]
    return start, end 

def indent_not_multiple_of_tab_size(project_folder, tab_size, *args, **kwargs):
    """
        Since there are cases for which col_offset is computed incorrectly,
        this validator must be nothing more than a simple warning.

        It compliments the pep8 validator which tends to fail in cases when
        the indent is incorrect.
    """
    node_types_to_validate = (ast.For, ast.If, ast.FunctionDef, ast.With)
    for parsed_file in project_folder.get_parsed_py_files():
        lines_offsets = file_helpers.get_line_offsets(parsed_file.content)
        for node in ast.walk(parsed_file.ast_tree):
            if not ast_helpers.is_node_offset_fine(
                node,
                lines_offsets,
                node_types_to_validate,
                tab_size,
            ):
                return parsed_file.get_name_with_line(node.lineno) 

def _get_modules(self, content) -> Set[str]:
        imports = set()
        tree = ast.parse(content)
        for node in ast.walk(tree):
            if isinstance(node, ast.Import):
                for subnode in node.names:
                    imports.add(subnode.name)
            elif isinstance(node, ast.ImportFrom) and node.level == 0:
                imports.add(node.module)
        modules = set()
        for module in imports:
            if not module:
                continue
            module = module.split('.', maxsplit=1)[0]
            if module in self.stdlib:
                continue
            module = self.aliases.get(module, module)
            modules.add(module)
        return modules 

def get_statement_startend2(lineno, node):
    import ast

    # flatten all statements and except handlers into one lineno-list
    # AST's line numbers start indexing at 1
    values = []
    for x in ast.walk(node):
        if isinstance(x, (ast.stmt, ast.ExceptHandler)):
            values.append(x.lineno - 1)
            for name in ("finalbody", "orelse"):
                val = getattr(x, name, None)
                if val:
                    # treat the finally/orelse part as its own statement
                    values.append(val[0].lineno - 1 - 1)
    values.sort()
    insert_index = bisect_right(values, lineno)
    start = values[insert_index - 1]
    if insert_index >= len(values):
        end = None
    else:
        end = values[insert_index]
    return start, end 

def get_statement_startend2(lineno, node):
    import ast
    # flatten all statements and except handlers into one lineno-list
    # AST's line numbers start indexing at 1
    l = []
    for x in ast.walk(node):
        if isinstance(x, _ast.stmt) or isinstance(x, _ast.ExceptHandler):
            l.append(x.lineno - 1)
            for name in "finalbody", "orelse":
                val = getattr(x, name, None)
                if val:
                    # treat the finally/orelse part as its own statement
                    l.append(val[0].lineno - 1 - 1)
    l.sort()
    insert_index = bisect_right(l, lineno)
    start = l[insert_index - 1]
    if insert_index >= len(l):
        end = None
    else:
        end = l[insert_index]
    return start, end 

def get_statement_startend2(lineno, node):
    import ast

    # flatten all statements and except handlers into one lineno-list
    # AST's line numbers start indexing at 1
    values = []
    for x in ast.walk(node):
        if isinstance(x, (ast.stmt, ast.ExceptHandler)):
            values.append(x.lineno - 1)
            for name in ("finalbody", "orelse"):
                val = getattr(x, name, None)
                if val:
                    # treat the finally/orelse part as its own statement
                    values.append(val[0].lineno - 1 - 1)
    values.sort()
    insert_index = bisect_right(values, lineno)
    start = values[insert_index - 1]
    if insert_index >= len(values):
        end = None
    else:
        end = values[insert_index]
    return start, end 

def get_statement_startend2(lineno, node):
    import ast
    # flatten all statements and except handlers into one lineno-list
    # AST's line numbers start indexing at 1
    l = []
    for x in ast.walk(node):
        if isinstance(x, _ast.stmt) or isinstance(x, _ast.ExceptHandler):
            l.append(x.lineno - 1)
            for name in "finalbody", "orelse":
                val = getattr(x, name, None)
                if val:
                    # treat the finally/orelse part as its own statement
                    l.append(val[0].lineno - 1 - 1)
    l.sort()
    insert_index = bisect_right(l, lineno)
    start = l[insert_index - 1]
    if insert_index >= len(l):
        end = None
    else:
        end = l[insert_index]
    return start, end 

def snoop(self, _line, cell):
        filename = self.shell.compile.cache(cell)
        code = self.shell.compile(cell, filename, 'exec')
        tracer = snoop()

        tracer.variable_whitelist = set()
        for node in ast.walk(ast.parse(cell)):
            if isinstance(node, ast.Name):
                name = node.id
                
                if isinstance(
                        self.shell.user_global_ns.get(name),
                        type(ast),  
                ):
                    # hide modules
                    continue
                    
                tracer.variable_whitelist.add(name)

        tracer.target_codes.add(code)
        with tracer:
            self.shell.ex(code) 

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 ast_names(code):
    """Iterator that yields all the (ast) names in a Python expression.

    :arg code: A string containing a Python expression.
    """
    # Syntax that allows new name bindings to be introduced is tricky to
    # handle here, so we just refuse to do so.
    disallowed_ast_nodes = (ast.Lambda, ast.ListComp, ast.GeneratorExp)
    if sys.version_info >= (2, 7):
        disallowed_ast_nodes += (ast.DictComp, ast.SetComp)

    for node in ast.walk(ast.parse(code)):
        if isinstance(node, disallowed_ast_nodes):
            raise PatsyError("Lambda, list/dict/set comprehension, generator "
                             "expression in patsy formula not currently supported.")
        if isinstance(node, ast.Name):
            yield node.id 

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 walk_python_files():
    u'''
    Generator that yields all CKAN Python source files.

    Yields 2-tuples containing the filename in absolute and relative (to
    the project root) form.
    '''
    def _is_dir_ignored(root, d):
        if d.startswith(u'.'):
            return True
        return os.path.join(rel_root, d) in IGNORED_DIRS

    for abs_root, dirnames, filenames in os.walk(PROJECT_ROOT):
        rel_root = os.path.relpath(abs_root, PROJECT_ROOT)
        if rel_root == u'.':
            rel_root = u''
        dirnames[:] = [d for d in dirnames if not _is_dir_ignored(rel_root, d)]
        for filename in filenames:
            if not filename.endswith(u'.py'):
                continue
            abs_name = os.path.join(abs_root, filename)
            rel_name = os.path.join(rel_root, filename)
            yield abs_name, rel_name 

def get_version(fname):
    """Get the version info from the mpld3 package without importing it"""
    import ast

    with open(fname) as init_file:
        module = ast.parse(init_file.read())

    version = (ast.literal_eval(node.value) for node in ast.walk(module)
               if isinstance(node, ast.Assign)
               and node.targets[0].id == "__version__")
    try:
        return next(version)
    except StopIteration:
        raise ValueError("version could not be located")


#cffi.verifier.cleanup_tmpdir() 

def ingest(self, rootdir):
        """
        Collect all the .py files to perform analysis upon
        """
        if not os.path.isdir(rootdir):
            raise Exception("directory %s passed in is not a dir" % rootdir)
        
        self.__target_dir = rootdir 

        # walk the dirs/files
        for root, subdir, files in os.walk(self.__target_dir):
            for f in files:
                if f.endswith(".py"):
                    fullpath = root + os.sep + f
                    contents = file(fullpath).read()
                    tree = ast.parse(contents)
                    self.__fn_to_ast[fullpath] = tree 

        # potentially analyze .html files for jinja templates
        if self.perform_jinja_analysis:
            self.__template_dir = self.get_template_dir() 

def find_template_dir(self):
        # web-p2 is web-p2/partners/templates
        # login is login/templates
        # TODO: look for invocations of `jinja2.Environment` and see if
        # we can pull the template directory / package from there? Should work
        # for most.
        template_dirs = set()
        for root, subdir, files in os.walk(self.__target_dir):
            for fname in files:
                fpath = os.path.join(root, fname)
                if fname.endswith(".html"):
                    with open(fpath, "rb") as f:
                        # Hmm, smells like a jinja template!
                        if b"{%" in f.read():
                            template_dirs.add(root)
        # If there are multiple template directories in a repo we might need
        # repo-specific overrides.
        return None if not template_dirs else os.path.commonprefix(template_dirs) 

def get_template_dir(self):
        """ return the directory containing jinja2 templates
            ex:  web-p2 is web-p2/partners/templates
        """
        template_dirs = set()
        for root, subdir, files in os.walk(self.__target_dir):
            for fname in files:
                fpath = os.path.join(root, fname)
                if fname.endswith(".html"):
                    with open(fpath, "rb") as f:
                        # Hmm, smells like a jinja template!
                        if b"{%" in f.read():
                            template_dirs.add(root)
        # If there are multiple template directories in a repo we might need
        # repo-specific overrides.
        return None if not template_dirs else os.path.commonprefix(template_dirs) 

def _fine_property_definition(self, property_name):
        """Find the lines in the source code that contain this property's name and definition.

        This function can find both attribute assignments as well as methods/functions.

        Args:
            property_name (str): the name of the property to look up in the template definition

        Returns:
            tuple: line numbers for the start and end of the attribute definition
        """
        for node in ast.walk(ast.parse(self._source)):
            if isinstance(node, ast.Assign) and node.targets[0].id == property_name:
                return node.targets[0].lineno - 1, self._get_node_line_end(node)
            elif isinstance(node, ast.FunctionDef) and node.name == property_name:
                return node.lineno - 1, self._get_node_line_end(node)
        raise ValueError('The requested node could not be found.') 

def get_statement_startend2(lineno: int, node: ast.AST) -> Tuple[int, Optional[int]]:
    # flatten all statements and except handlers into one lineno-list
    # AST's line numbers start indexing at 1
    values = []  # type: List[int]
    for x in ast.walk(node):
        if isinstance(x, (ast.stmt, ast.ExceptHandler)):
            values.append(x.lineno - 1)
            for name in ("finalbody", "orelse"):
                val = getattr(x, name, None)  # type: Optional[List[ast.stmt]]
                if val:
                    # treat the finally/orelse part as its own statement
                    values.append(val[0].lineno - 1 - 1)
    values.sort()
    insert_index = bisect_right(values, lineno)
    start = values[insert_index - 1]
    if insert_index >= len(values):
        end = None
    else:
        end = values[insert_index]
    return start, end 

def _parse_setup_params(setup_py: Path) -> Dict[str, Any]:
    with setup_py.open("r", encoding="utf8") as sp:
        setup_tree = ast.parse(sp.read())

    LOG.debug(f"AST visiting {setup_py}")
    for node in ast.walk(setup_tree):
        if isinstance(node, ast.Assign):
            for target in node.targets:
                target_id = getattr(target, "id", None)
                if not target_id:
                    continue

                if target_id == "ptr_params":
                    LOG.debug(f"Found ptr_params in {setup_py}")
                    return dict(ast.literal_eval(node.value))
    return {} 

def read_tags(filename):
    """Reads values of "magic tags" defined in the given Python file.

    :param filename: Python filename to read the tags from
    :return: Dictionary of tags
    """
    with open(filename) as f:
        ast_tree = ast.parse(f.read(), filename)

    res = {}
    for node in ast.walk(ast_tree):
        if type(node) is not ast.Assign:
            continue

        target = node.targets[0]
        if type(target) is not ast.Name:
            continue
        if not (target.id.startswith('__') and target.id.endswith('__')):
            continue

        name = target.id[2:-2]
        res[name] = ast.literal_eval(node.value)

    return res 

def linerange(node):
    """Get line number range from a node."""
    strip = {"body": None, "orelse": None,
             "handlers": None, "finalbody": None}
    for key in strip.keys():
        if hasattr(node, key):
            strip[key] = getattr(node, key)
            node.key = []

    lines_min = 9999999999
    lines_max = -1
    for n in ast.walk(node):
        if hasattr(n, 'lineno'):
            lines_min = min(lines_min, n.lineno)
            lines_max = max(lines_max, n.lineno)

    for key in strip.keys():
        if strip[key] is not None:
            node.key = strip[key]

    if lines_max > -1:
        return list(range(lines_min, lines_max + 1))
    return [0, 1] 

def setup_py(location: str) -> dict:
    """Returns back any configuration info we are able to determine from a setup.py file"""
    setup_config = {}
    try:
        with open(location) as setup_py_file:
            for node in ast.walk(ast.parse(setup_py_file.read())):
                if (
                    type(node) == _ast.Call
                    and type(getattr(node, "func", None)) == _ast.Name
                    and node.func.id == "setup"  # type: ignore
                ):
                    for keyword in node.keywords:  # type: ignore
                        if keyword.arg == "packages":
                            setup_config["modules"] = ast.literal_eval(keyword.value)
                            break
                    break
    except Exception as error:
        warnings.warn(f"Error ({error}) occurred trying to parse setup.py file: {location}")

    return setup_config 

def get_statement_startend2(lineno, node):
    import ast
    # flatten all statements and except handlers into one lineno-list
    # AST's line numbers start indexing at 1
    l = []
    for x in ast.walk(node):
        if isinstance(x, _ast.stmt) or isinstance(x, _ast.ExceptHandler):
            l.append(x.lineno - 1)
            for name in "finalbody", "orelse":
                val = getattr(x, name, None)
                if val:
                    # treat the finally/orelse part as its own statement
                    l.append(val[0].lineno - 1 - 1)
    l.sort()
    insert_index = bisect_right(l, lineno)
    start = l[insert_index - 1]
    if insert_index >= len(l):
        end = None
    else:
        end = l[insert_index]
    return start, end 

def get_top_imported_names(file: str) -> Set[str]:
    """Collect names imported in given file.

    We only collect top-level names, i.e. `from foo.bar import baz`
    will only add `foo` to the list.
    """
    if not file.endswith(".pyi"):
        return set()
    with open(os.path.join(file), "rb") as f:
        content = f.read()
    parsed = ast.parse(content)
    top_imported = set()
    for node in ast.walk(parsed):
        if isinstance(node, ast.Import):
            for name in node.names:
                top_imported.add(name.name.split('.')[0])
        elif isinstance(node, ast.ImportFrom):
            if node.level > 0:
                # Relative imports always refer to the current package.
                continue
            assert node.module
            top_imported.add(node.module.split('.')[0])
    return top_imported 

def iter_type_comments(self):
        """
        This feature is only available for python 3.8.

        PEP 526 -- Syntax for Variable Annotations
        https://www.python.org/dev/peps/pep-0526/

        https://docs.python.org/3.8/library/ast.html#ast.parse

        """
        buffer = io.StringIO(self.source)
        for token in tokenize.generate_tokens(buffer.readline):
            if token.type == tokenize.COMMENT:
                comment_string = token.string.split("# type: ")
                if comment_string != [token.string]:
                    try:
                        functype = ast.parse(
                            comment_string[1], mode="func_type"
                        )
                    except SyntaxError as err:
                        if self.show_error:
                            error_messages = f"{token.line}\n{comment_string[1]} {Color(str(err)).red}"
                            print(error_messages)
                    else:
                        for node in ast.walk(
                            ast.Module(functype.argtypes + [functype.returns])
                        ):
                            if isinstance(node, ast.Name) and isinstance(
                                node.ctx, ast.Load
                            ):
                                yield node 

def visit_ClassDef(self, node):
        for function_node in ast.walk(node):
            if isinstance(function_node, ast.FunctionDef):
                function_node.class_def = True
        self.classes.append({"lineno": node.lineno, "name": node.name}) 

def visit_Attribute(self, node):
        local_attr = []
        for attr_node in ast.walk(node):
            if isinstance(attr_node, ast.Name):
                local_attr.append(attr_node.id)
            elif isinstance(attr_node, ast.Attribute):
                local_attr.append(attr_node.attr)
        local_attr.reverse()
        self.names.append(
            {"lineno": node.lineno, "name": ".".join(local_attr)}
        )