Question: How to walk/visit the AST?

[Peter554]

I'm trying to extract all the imports from a python file. To do this I'm using the rustpython_parser::parse function to get an AST and then iterating over all the statements in the body to find Stmt::Import and Stmt::ImportFrom. This works for imports defined at the root of the file, however misses any imports defined inside e.g. function/class definitions etc. So I think what I really want to do is walk the AST and visit all the nodes. Is there some way to achieve this with one of the rustpython packages?

[youknowone]

Hi, rustpython_ast crate contains a few utilities like Visitor and Fold. They must be helpful.

[bnkc]

Hey there. I found the visitor.rs module but cannot import it. are there any updates?

[fanninpm]

Is the module public?

[bnkc]

I'm decently new to rust but from my understanding it is:

#[cfg(feature = "visitor")]
pub use visitor::Visitor;

I've also enabled the feature @fanninpm

[fanninpm]

The following line likely needs to be changed to say pub mod visitor {:

Parser/ast/src/lib.rs

Line 46 in 9ce55ae

mod visitor {

[bnkc]

Ok should i open a PR?

[fanninpm]

Give it a shot!

[youknowone]

https://github.com/RustPython/Parser/blob/9ce55aefdeb35e2f706ce0b02d5a2dfe6295fc57/ast/src/lib.rs#L60C1-L60C26

Isn't this line allow to import rustpython_ast::Visitor?

[miguelcsilva]

I'm very new to rust and was also trying to walk the AST in order to find all usages of a particular function.
Is that functionality provided by this module? Any example of out to do this would be amazing.

[Ivanmatthew]

Hi all,

I'll write down a brief explanation of how to add the visitor and how to utilize it.
You'll have to install the packages rustpython-parser for the actual parsing and rustpython-ast for interacting with the AST.
When adding rustpython-ast, you have to add the feature "visitor".
It will look something like the following, in Cargo.toml (root of your project)

[dependencies]
rustpython-ast = { version = "0.4.0", features = ["visitor"] }
rustpython-parser = "0.4.0"

Once you did that, you can use the crate like so:

use rustpython_ast::Visitor;

Visitor is a public trait, which means you'll have to create an implementation.
The easiest way to interact with the visitor is to create your own struct with properties you'd like to interact with when traversing the AST.
For example keeping a count of attributes are in the script.
It will look something like so:

struct AttributeCounter {
    attributes_count: usize
}

// Here you implement the trait for a struct
impl Visitor for AttributeCounter {
    // We'll get to this later
}

Nice!
Now by nature the currently implemented methods are inherited from the trait. Meaning all the methods defined in the Visitor trait will now be called.
You can override these functions by defining them yourself with the same signature (same parameters and name).

The Visitor is built to be a recursive walker/visitor, starting with the function visit_stmt.
From thereon it calls it's generic_visit_stmt function which then walks through every node present in the stmt node, branching out to other AST nodes like if, while, for, function def, attribute etc.
The easiest way to start counting attributes would be to override the method that's visiting attributes, this is visit_expr_attribute at the time of writing.

impl Visitor for AttributeCounter {
    fn visit_expr_attribute(&mut self, node: ExprAttribute) {
        self.attributes_count += 1;
    }
}

Great, it now counts attributes!
However, we missed something here, because we have overridden the original implementation of the function, it will not traverse further down the node and stops here.
If we look at the source code of the visitor, it is calling to a generic implementation which drills down to the other nodes. So to ensure the whole AST is traversed, we should also keep the original implementation in our code. Which becomes:

impl Visitor for AttributeCounter {
    fn visit_expr_attribute(&mut self, node: ExprAttribute) {
        self.attributes_count += 1;
        
        self.generic_visit_expr_attribute(node);
    }
}

Now the visitor will visit the whole AST and count any attributes for you.