Walking a python AST

1

I am building a CLI app that needs to walk the python AST. I ended up using the following crate:
rustpython-parser

here's my code:

fn visit_ast(nodes: &[Stmt], collected_deps: &mut HashSet<String>) {
    nodes.iter().for_each(|node| match node {
        Stmt::Import(import) => {
            import.names.iter().for_each(|alias| {
                collected_deps.insert(stem_import(&alias.name));
            });
        }
        Stmt::ImportFrom(import) => {
            if let Some(module) = &import.module {
                collected_deps.insert(stem_import(module));
            }
        }
        Stmt::FunctionDef(function_def) => visit_ast(&function_def.body, collected_deps),
        Stmt::ClassDef(class_def) => visit_ast(&class_def.body, collected_deps),
        Stmt::AsyncFunctionDef(async_function_def) => {
            visit_ast(&async_function_def.body, collected_deps)
        }
        Stmt::For(for_stmt) => visit_ast(&for_stmt.body, collected_deps),
        Stmt::AsyncFor(async_for_stmt) => visit_ast(&async_for_stmt.body, collected_deps),
        Stmt::While(while_stmt) => visit_ast(&while_stmt.body, collected_deps),
        Stmt::If(if_stmt) => visit_ast(&if_stmt.body, collected_deps),
        Stmt::With(with_stmt) => visit_ast(&with_stmt.body, collected_deps),
        Stmt::AsyncWith(async_with_stmt) => visit_ast(&async_with_stmt.body, collected_deps),
        Stmt::Match(match_stmt) => {
            match_stmt.cases.iter().for_each(|case| {
                visit_ast(&case.body, collected_deps);
            });
        }
        Stmt::Try(try_stmt) => {
            visit_ast(&try_stmt.body, collected_deps);
            visit_ast(&try_stmt.orelse, collected_deps);
            visit_ast(&try_stmt.finalbody, collected_deps);
        }
        Stmt::TryStar(try_star_stmt) => {
            visit_ast(&try_star_stmt.body, collected_deps);
            visit_ast(&try_star_stmt.orelse, collected_deps);
            visit_ast(&try_star_stmt.finalbody, collected_deps);
        }

        _ => {}
    });
}

pub fn get_used_imports(dir: &Path) -> Result<HashSet<String>> {
    WalkDir::new(dir)
        .into_iter()
        .filter_map(|e| e.ok())
        .filter(|entry| entry.file_name().to_string_lossy().ends_with(".py"))
        .try_fold(HashSet::new(), |mut acc, entry| {
            let file_content = fs::read_to_string(entry.path())?;
            let module = parse(&file_content, Mode::Module, "<embedded>")?;
            let nodes = &module.module().unwrap().body; // Maybe should do a match here??
            let mut collected_deps: HashSet<String> = HashSet::new();

            visit_ast(&nodes, &mut collected_deps);

            acc.extend(collected_deps);

            Ok(acc)
        })
}

The Goal is to collect all imports that may be nested in various namespaces.

I'm pretty new to rust but I was surprised the crate Had a visitor pattern but didn't make it public. When I walk really large trees, this can take longer than I would like. Does anyone have experience with highly recursive tasks like this? Any suggestions of how to improve this logic?

2

The Visitor trait is public with visitor feature enabled in rustpython-ast crate. lib.rs - source

rustpython-parser doesn't enable visitor, rustpython-ast neither by default.

4

Sweet thanks!

5

Just want to confirm - are you running with --release?

6

@mr_rustbot I am not, not yet atleast. I wanted to get a little further along in the project. Is the resulting binary much faster?

7

Yes, the default build profile is debug and is extremely slow compared to --release.

8

This topic was automatically closed 90 days after the last reply. We invite you to open a new topic if you have further questions or comments.