Import discovery

Cargo.toml

@@ -8,6 +8,5 @@ edition = "2021"
 [dependencies]
 anyhow = "1.0.75"
 rayon = "1.8.0"
-serde = { version="1.0.192", features = ["derive", "rc"] }
-serde_json = "1.0.108"
+rustpython-parser = "0.3.0"
 thiserror = "1.0.50"

example/__init__.py

@@ -0,0 +1,25 @@
+import datetime
+from os import path
+
+import example.a
+import example.child.c_a
+
+def foo():
+    from example import b
+    from example.child import c_b
+
+class Bar:
+    from example.c import C
+    from example.child.c_c import C
+
+from . import d
+from .child import c_d
+
+from .e import E
+from .child.c_e import E
+
+import example.child
+from example import child2
+from example.child3 import CHILD
+from . import child4
+from .child5 import CHILD

example/a.py

@@ -0,0 +1 @@
+A = "A"

example/c.py

@@ -0,0 +1 @@
+C = "C"

example/child/__init__.py

@@ -0,0 +1,23 @@
+import datetime
+from os import path
+
+import example.a
+import example.child.c_a
+
+from example import b
+from example.child import c_b
+
+from example.c import C
+from example.child.c_c import C
+
+from .. import d
+from . import c_d
+
+from ..e import E
+from .c_e import E
+
+import example
+from example import child2
+from example.child3 import CHILD
+from .. import child4
+from ..child5 import CHILD

example/child/c_c.py

@@ -0,0 +1 @@
+C = "C"

example/child/c_d.py

@@ -0,0 +1 @@
+D = "D"

example/child/c_e.py

@@ -0,0 +1 @@
+E = "E"

example/child/c_z.py

@@ -0,0 +1,23 @@
+import datetime
+from os import path
+
+import example.a
+import example.child.c_a
+
+from example import b
+from example.child import c_b
+
+from example.c import C
+from example.child.c_c import C
+
+from .. import d
+from . import c_d
+
+from ..e import E
+from .c_e import E
+
+from .. import example
+from example import child2
+from example.child3 import CHILD
+from .. import child4
+from ..child5 import CHILD

example/child3/__init__.py

@@ -0,0 +1 @@
+CHILD = "CHILD"

example/child5/__init__.py

@@ -0,0 +1 @@
+CHILD = "CHILD"

example/d.py

@@ -0,0 +1 @@
+D = "D"

example/e.py

@@ -0,0 +1 @@
+E = "E"

example/z.py

@@ -0,0 +1,23 @@
+import datetime
+from os import path
+
+import example.a
+import example.child.c_a
+
+from example import b
+from example.child import c_b
+
+from example.c import C
+from example.child.c_c import C
+
+from . import d
+from .child import c_d
+
+from .e import E
+from .child.c_e import E
+
+import example.child
+from example import child2
+from example.child3 import CHILD
+from . import child4
+from .child5 import CHILD

src/ast_visit.rs

@@ -0,0 +1,189 @@
+use rustpython_parser::ast::{ExceptHandler, ModModule, Stmt};
+
+pub trait StatementVisitor {
+    fn visit(&mut self, stmt: &Stmt) -> bool;
+}
+
+pub fn visit_statements<T>(ast: &ModModule, visitor: &mut T)
+where
+    T: StatementVisitor,
+{
+    for stmt in ast.body.iter() {
+        visit_stmt(stmt, visitor);
+    }
+}
+
+fn visit_stmt<T>(stmt: &Stmt, visitor: &mut T)
+where
+    T: StatementVisitor,
+{
+    let ok = visitor.visit(stmt);
+    if !ok {
+        return;
+    }
+    match stmt {
+        Stmt::FunctionDef(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::AsyncFunctionDef(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::ClassDef(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::Return(_) => return,
+        Stmt::Delete(_) => return,
+        Stmt::Assign(_) => return,
+        Stmt::TypeAlias(_) => return,
+        Stmt::AugAssign(_) => return,
+        Stmt::AnnAssign(_) => return,
+        Stmt::For(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+            for stmt in def.orelse.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::AsyncFor(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+            for stmt in def.orelse.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::While(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+            for stmt in def.orelse.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::If(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+            for stmt in def.orelse.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::With(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::AsyncWith(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::Match(def) => {
+            for case in def.cases.iter() {
+                for stmt in case.body.iter() {
+                    visit_stmt(&stmt, visitor);
+                }
+            }
+        }
+        Stmt::Raise(_) => return,
+        Stmt::Try(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+            for handler in def.handlers.iter() {
+                match handler {
+                    ExceptHandler::ExceptHandler(handler) => {
+                        for stmt in handler.body.iter() {
+                            visit_stmt(&stmt, visitor);
+                        }
+                    }
+                }
+            }
+            for stmt in def.orelse.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+            for stmt in def.finalbody.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::TryStar(def) => {
+            for stmt in def.body.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+            for handler in def.handlers.iter() {
+                match handler {
+                    ExceptHandler::ExceptHandler(handler) => {
+                        for stmt in handler.body.iter() {
+                            visit_stmt(&stmt, visitor);
+                        }
+                    }
+                }
+            }
+            for stmt in def.orelse.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+            for stmt in def.finalbody.iter() {
+                visit_stmt(&stmt, visitor);
+            }
+        }
+        Stmt::Assert(_) => return,
+        Stmt::Import(_) => return,
+        Stmt::ImportFrom(_) => return,
+        Stmt::Global(_) => return,
+        Stmt::Nonlocal(_) => return,
+        Stmt::Expr(_) => return,
+        Stmt::Pass(_) => return,
+        Stmt::Break(_) => return,
+        Stmt::Continue(_) => return,
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use rustpython_parser::ast::Mod;
+    use std::{fs, path::Path};
+
+    use super::*;
+
+    struct TestVisitor {
+        counter: u64,
+    }
+
+    impl StatementVisitor for TestVisitor {
+        fn visit(&mut self, stmt: &Stmt) -> bool {
+            self.counter += 1;
+            return true;
+        }
+    }
+
+    #[test]
+    fn test_visit() {
+        let path = Path::new("./src/ast_visit_test.py");
+        let code = fs::read_to_string(path).unwrap();
+        let parse_result = rustpython_parser::parse(
+            &code,
+            rustpython_parser::Mode::Module,
+            path.to_str().unwrap(),
+        );
+        let ast = match parse_result {
+            Ok(m) => match m {
+                Mod::Module(mm) => mm,
+                _ => panic!(),
+            },
+            _ => panic!(),
+        };
+
+        let mut visitor = TestVisitor { counter: 0 };
+
+        visit_statements(&ast, &mut visitor);
+
+        assert_eq!(visitor.counter, 5);
+    }
+}

src/ast_visit_test.py

@@ -0,0 +1,7 @@
+foo = "foo"
+
+class A:
+    bar = "bar"
+
+def b():
+    baz = "baz"