trying to implement the root dir check

Author: deadlovelll

src/core/ast_parser/ast_parser.rs

@@ -32,7 +32,7 @@ impl AstParser {
         }
     }
 
-    pub fn parse(&self, path: &Path) {
+    pub fn parse(&self, path: &Path, root_dirs: &Vec<String>) {
         let mut variables: HashMap<String, Vec<String>> = HashMap::new();
         let code = fs::read_to_string(path).expect("Failed to read file");
         let language: Language = tree_sitter_python::LANGUAGE.into();
@@ -45,6 +45,7 @@ impl AstParser {
             &code, 
             &query, 
             root_node,
+            root_dirs,
         );
         let imports = result.0;
         let defines = result.1;

src/core/ast_parser/import_classifier.rs

@@ -1,49 +1,12 @@
-use std::process::Command;
-use std::path::Path;
-
 /// Classifies Python imports as standard library or third-party/project-specific.
 pub struct ImportClassifier {
-    /// List of standard library module names.
-    stdlib_modules: Vec<String>
 }
 
 impl ImportClassifier {
     /// Creates a new `ImportClassifier` and initializes the
     /// list of standard library modules.
     pub fn new() -> Self {
-        let stdlib_modules: Vec<String> = [
-            "__phello__", "_pyrepl", "asyncio", "collections", "compression",
-            "concurrent", "ctypes", "curses", "dbm", "email", "encodings",
-            "ensurepip", "html", "http", "idlelib", "importlib", "json",
-            "logging", "multiprocessing", "pathlib", "profiling", "pydoc_data",
-            "re", "site-packages", "sqlite3", "string", "sysconfig", "test",
-            "tkinter", "tomllib", "turtledemo", "unittest", "urllib", "venv",
-            "wsgiref", "xml", "xmlrpc", "zipfile", "zoneinfo", "__future__", "__hello__",
-            "_aix_support", "_android_support", "_apple_support", "_ast_unparse",
-            "_collections_abc", "_colorize", "_compat_pickle", "_ios_support", "_markupbase",
-            "_opcode_metadata", "_osx_support", "_py_abc", "_py_warnings", "_pydatetime",
-            "_pydecimal", "_pyio", "_pylong", "_sitebuiltins", "_strptime", "_threading_local",
-            "_weakrefset", "abc", "annotationlib", "antigravity", "argparse", "ast", "base64",
-            "bdb", "bisect", "bz2", "cProfile", "calendar", "cmd", "code", "codecs",
-            "codeop", "colorsys", "compileall", "configparser", "contextlib", "contextvars",
-            "copy", "copyreg", "csv", "dataclasses", "datetime", "decimal", "difflib",
-            "dis", "doctest", "enum", "filecmp", "fileinput", "fnmatch", "fractions",
-            "ftplib", "functools", "genericpath", "getopt", "getpass", "gettext", "glob",
-            "graphlib", "gzip", "hashlib", "heapq", "hmac", "imaplib", "inspect", "io",
-            "ipaddress", "keyword", "linecache", "locale", "lzma", "mailbox", "mimetypes",
-            "modulefinder", "netrc", "ntpath", "nturl2path", "numbers", "opcode", "operator",
-            "optparse", "os", "pdb", "pickle", "pickletools", "pkgutil", "platform", "plistlib",
-            "poplib", "posixpath", "pprint", "profile", "pstats", "pty", "py_compile", "pyclbr",
-            "pydoc", "queue", "quopri", "random", "reprlib", "rlcompleter", "runpy", "sched",
-            "secrets", "selectors", "shelve", "shlex", "shutil", "signal", "site", "smtplib",
-            "socket", "socketserver", "ssl", "stat", "statistics", "stringprep", "struct",
-            "subprocess", "symtable", "tabnanny", "tarfile", "tempfile", "textwrap", "this",
-            "threading", "timeit", "token", "tokenize", "trace", "traceback", "tracemalloc",
-            "tty", "turtle", "types", "typing", "uuid", "warnings", "wave", "weakref",
-            "webbrowser", "zipapp", "zipimport"
-        ].iter().map(|s| s.to_string()).collect();
-
-        Self {stdlib_modules}
+        Self {}
     }
 
     /// Returns `true` if the import is project level.
@@ -59,38 +22,17 @@ impl ImportClassifier {
     /// assert_eq!(classifier.is_eligible(&"infra.db.modules".to_string()), true);
     /// assert_eq!(classifier.is_eligible(&"os".to_string()), false);
     /// ```
-    pub fn is_eligible(&self, import: &String) -> bool {
-        // Checks is import from standard
-        // library or not
-        true
-        // if self.stdlib_modules.contains(import) {
-        //     return false;
-        // } else if self.is_third_party(import) {
-        //     return false;
-        // } else {
-        //     return true;
-        // }
-    }
+    pub fn is_eligible(
+        &self, 
+        import: &String,
+        root_dirs: &Vec<String>,
+    ) -> bool {
 
-    // fn is_third_party(&self, import: &str) -> bool {
-    //     let site_packages = self.get_site_packages();
-    //     for sp in site_packages {
-    //         let path = Path::new(&sp).join(import);
-    //         println!("{}", path.to_string_lossy());
-    //         if path.exists() {
-    //             return true;
-    //         }
-    //     }
-    //     false
-    // }
-
-    // fn get_site_packages(&self) -> Vec<String> {
-    //     let output = Command::new("python3")
-    //         .args(&["-c", "import site; print('\\n'.join(site.getsitepackages()))"])
-    //         .output()
-    //         .expect("Failed to run python3");
-
-    //     let stdout = String::from_utf8_lossy(&output.stdout);
-    //     stdout.lines().map(|s| s.to_string()).collect()
-    // }
+        for dir in root_dirs {
+            if import.starts_with(dir) {
+                return true;
+            }
+        }
+        return false;
+    }
 }

src/core/ast_parser/tree_analyzer.rs

@@ -1,5 +1,3 @@
-use std::fmt::format;
-
 use tree_sitter::{Query, Node, QueryCursor, StreamingIterator};
 
 use crate::core::ast_parser::import_classifier::ImportClassifier;
@@ -23,6 +21,7 @@ impl TreeAnalyzer {
         code: &str,
         query: &Query,
         root_node: Node,
+        root_dirs: &Vec<String>,
     ) -> (Vec<String>, Vec<String>, Vec<String>) {
 
         let mut imports: Vec<String> = vec![];
@@ -49,7 +48,7 @@ impl TreeAnalyzer {
                     }
                     "import" => {
                         let text = self.dot_name.get(&code, cap.node);
-                        let is_eligible: bool = self.import_classifier.is_eligible(&text);
+                        let is_eligible: bool = self.import_classifier.is_eligible(&text, root_dirs);
                         if is_eligible {
                             imports.push(text);
                         }
@@ -58,6 +57,7 @@ impl TreeAnalyzer {
                         let text = self.dot_name.get(&code, cap.node);
                         if !text.starts_with("__") & !text.ends_with("__") {
                             let mut parent = cap.node.parent();
+                            let mut found_class = false;
                             while let Some(p) = parent {
                                 if p.kind() == "class_definition" {
                                     if let Some(name_node) = p.child_by_field_name("name") {
@@ -66,17 +66,21 @@ impl TreeAnalyzer {
                                         .to_string();
                                         let class_method = format!("{class}.{text}");
                                         defines.push(class_method);
+                                        found_class = true;
                                     }
                                 }
                                 parent = p.parent();
                             }
+                            if !found_class {
+                                defines.push(text);
+                            }
                         }
                     }
                     _ => {}
                 }
             }
             if let Some(m) = module {
-                let is_eligible: bool = self.import_classifier.is_eligible(&m);
+                let is_eligible: bool = self.import_classifier.is_eligible(&m, root_dirs);
                 if is_eligible {
                     for n in names {
                         let full_import = format!("{}.{}", m, n);

src/core/file_walker/file_walker.rs

@@ -1,4 +1,5 @@
 use std::fs;
+use std::fs::DirEntry;
 use std::path::Path;
 
 use regex::Regex;
@@ -22,34 +23,58 @@ impl FileWalker {
     }
 
     pub fn run(&self, path: &Path) {
-        self.walk(path);
-        self.file_hasher.write_to_file();
+        let root_dirs = self.get_root_dirs(path);
+
+        self.walk(path, &root_dirs);
         self.graph_creator.write_to_file();
     }
 
-    pub fn walk(&self, path: &Path) {
-        let entries: Vec<_> = fs::read_dir(path)
-            .unwrap()
-            .map(|e| e.unwrap())
-            .collect();
-
+    pub fn walk(&self, path: &Path, root_dirs: &Vec<String>) {
+        let entries = self.get_entries(path);
         entries.par_iter().for_each(|entry| {
             if entry.file_type().unwrap().is_dir() {
-                self.walk(&entry.path());
+                self.walk(&entry.path(), root_dirs);
             } else if self.is_eligible(entry.path().to_str().unwrap()) {
                 if !entry.path().to_str().unwrap().contains("test") {
                     self.file_hasher.hash(&entry.path());
                 }
-                self.graph_creator.create_graph(&entry.path());
+                self.graph_creator.create_graph(&entry.path(), &root_dirs);
             }
         });
     }
 
+    fn get_entries(&self, path: &Path) -> Vec<DirEntry> {
+        let entries: Vec<_> = fs::read_dir(path)
+            .unwrap()
+            .map(|e| e.unwrap())
+            .collect();
+
+        return entries
+    }
+
     fn is_eligible(&self, path: &str) -> bool {
-        let ignore_regex = Regex::new(r"(^|/)(__+).*\.py$").unwrap();
-        return path.ends_with(".py") 
-        && !path.contains("venv")
-        && !path.contains("test")
-        && !ignore_regex.is_match(path)
+        let ignore_regex = Regex::new(
+            r"^(?!.*venv)(?!.*test)(?!.*(?:^|/)(__+).+\.py$).+\.py$"
+        ).unwrap();
+        return !ignore_regex.is_match(path)
+    }
+
+    fn is_eligible_dir(&self, path: &str) -> bool {
+        let dir_regex = Regex::new(r"^(?!.*cache)(?!__.*__)(?!\.venv).*").unwrap();
+        return dir_regex.is_match(path);
+    }
+
+    fn get_root_dirs(&self, path: &Path) -> Vec<String> {
+        let mut root_dirs = vec![];
+        let entries = self.get_entries(path);
+        for root_entry in entries {
+            if root_entry.path().is_dir() {
+                if self.is_eligible_dir(root_entry.path().to_str().unwrap()) {
+                    println!("{}", path.display().to_string());
+                    root_dirs.push(path.display().to_string());
+                }
+            }
+        }
+        return root_dirs
     }
 }

src/core/graph_creator/graph_creator.rs

@@ -13,8 +13,12 @@ impl GraphCreator {
         Self {ast_data: HashMap::new(), ast_parser: AstParser::new()}
     }
 
-    pub fn create_graph(&self, path: &Path) {
-        self.ast_parser.parse(path);
+    pub fn create_graph(
+        &self, 
+        path: &Path,
+        root_dirs: &Vec<String>,
+    ) {
+        self.ast_parser.parse(path, root_dirs);
     }
 
     pub fn write_to_file(&self) {

src/core/project_intializer/project_initializer.rs

@@ -49,7 +49,7 @@ impl ProjectInitializer {
     /// Returns an `io::Error` if configuration creation fails.
     pub fn init(&mut self) -> io::Result<()> {
         self.config_creator.create_config()?;
-        self.file_walker.run(Path::new("./"));
+        self.file_walker.run(Path::new("."));
         Ok(())
     }
 }