Add config support and parse imports

Author: 123vivekr

Cargo.toml

@@ -9,6 +9,7 @@ phf = { version = "0.11.1", features = ["macros"] }
 serde = { version = "1.0.152", features = ["derive"] }
 serde_json = "1.0.114"
 tree-sitter = "0.20.9"
+toml = "0.8.11"
 
 [build-dependencies]
 cc="*"

asterisk.toml

@@ -0,0 +1,30 @@
+[languages]
+  [languages.python]
+    [languages.python.matchers]
+      import_statement = "import_from_statement"
+      [languages.python.matchers.module_name]
+        field_name = "module_name"
+        kind = "dotted_name"
+
+      [languages.python.matchers.object_name]
+        field_name = "name"
+        kind = "dotted_name"
+
+      [languages.python.matchers.alias]
+        field_name = "alias"
+        kind = "identifier"
+
+  [languages.rust]
+    [languages.rust.matchers]
+      import_statement = "use_declaration"
+      [languages.rust.matchers.module_name]
+        field_name = "path"
+        kind = "identifier"
+
+      [languages.rust.matchers.object_name]
+        field_name = "name"
+        kind = "identifier"
+
+      [languages.rust.matchers.alias]
+        field_name = "alias"
+        kind = "identifier"

examples/cli.rs

@@ -1,9 +1,12 @@
+
 use asterisk::indexer::index_directory;
 use serde::Serialize;
 use serde_json::json;
 use std::env;
 use std::fs::File;
 use std::io::Write;
+use std::fs;
+use asterisk::config::Config;
 
 #[derive(Serialize)]
 struct Output {
@@ -12,9 +15,12 @@ struct Output {
 }
 
 fn main() {
+    let toml_str = fs::read_to_string("asterisk.toml").expect("Unable to read file");
+    let config = Config::from_toml(&toml_str).unwrap();
+
     let args: Vec<String> = env::args().collect();
     let dir_path = &args[1];
-    let (blocks, call_stack, call_graph) = index_directory(dir_path);
+    let (blocks, call_stack, call_graph) = index_directory(&config, dir_path);
 
     let output = Output { blocks, call_stack };
 

src/config.rs

@@ -0,0 +1,35 @@
+use serde::Deserialize;
+use std::collections::HashMap;
+
+#[derive(Deserialize, Debug)]
+pub struct Config {
+    pub languages: HashMap<String, Language>,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct Language {
+    pub matchers: Matchers,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct Matchers {
+    pub import_statement: String,
+    pub module_name: Matcher,
+    pub object_name: Matcher,
+    pub alias: Matcher,
+}
+
+#[derive(Deserialize, Debug)]
+pub struct Matcher {
+    pub field_name: String,
+    pub kind: String,
+}
+
+
+impl Config {
+    pub fn from_toml(toml_str: &str) -> Result<Self, toml::de::Error> {
+        let configs: Config = toml::from_str(toml_str).expect("Failed to parse TOML");
+
+        Ok(configs)
+    }
+}

src/indexer.rs

@@ -1,5 +1,6 @@
 use jwalk::WalkDir;
 use std::path::Path;
+use crate::config::Config;
 
 use crate::block::{Block, BlockType};
 use crate::call_graph::CallGraph;
@@ -30,7 +31,7 @@ pub fn generate_node_key(
     key
 }
 
-pub fn index_directory(dir_path: &str) -> (Vec<Block>, CallStack, CallGraph) {
+pub fn index_directory(config: &Config, dir_path: &str) -> (Vec<Block>, CallStack, CallGraph) {
     let mut blocks = Vec::new();
     let mut call_stack = CallStack::new();
 
@@ -40,7 +41,7 @@ pub fn index_directory(dir_path: &str) -> (Vec<Block>, CallStack, CallGraph) {
 
         if path.is_file() && is_supported_file(&path) {
             let module_name = path.to_str().unwrap();
-            let file_blocks = parse_file(&path, module_name);
+            let file_blocks = parse_file(&path, module_name, &config);
             blocks.extend(file_blocks.clone());
 
             for block in &file_blocks {

src/lib.rs

@@ -4,3 +4,4 @@ pub mod call_stack;
 pub mod indexer;
 pub mod parser;
 pub mod utils;
+pub mod config;

src/parser.rs

@@ -1,4 +1,6 @@
 use crate::block::{Block, BlockType};
+use crate::config::{Config, Matchers};
+use jwalk::rayon::str::ParallelString;
 use std::collections::{HashMap, HashSet};
 use std::fs;
 use std::path::Path;
@@ -12,7 +14,7 @@ extern "C" {
     // Add more language bindings here
 }
 
-pub fn parse_file(file_path: &Path, module_name: &str) -> Vec<Block> {
+pub fn parse_file(file_path: &Path, module_name: &str, config: &Config) -> Vec<Block> {
     let code = fs::read_to_string(file_path).unwrap();
     let language = tree_sitter_language(file_path);
     let mut parser = Parser::new();
@@ -33,6 +35,7 @@ pub fn parse_file(file_path: &Path, module_name: &str) -> Vec<Block> {
         None,
         module_name,
         &mut imports,
+        &config,
     );
 
     if !non_function_blocks.is_empty() {
@@ -71,12 +74,14 @@ fn traverse_tree(
     class_name: Option<String>,
     module_name: &str,
     imports: &mut HashMap<String, String>,
+    config: &Config,
 ) {
     let node = cursor.node();
     let kind = node.kind();
 
     if is_import_statement(kind, language) {
-        if let Some((module, alias)) = parse_import_statement(code, node, language) {
+        if let Some((module, alias)) = parse_import_statement(code, node, language, config) {
+            println!("Module: {}, Alias: {}", module, alias);
             imports.insert(alias, module);
         }
     } else if is_class_definition(kind, language) {
@@ -98,6 +103,7 @@ fn traverse_tree(
                         Some(extracted_class_name.clone()),
                         module_name,
                         imports,
+                        config,
                     );
                     if !cursor.goto_next_sibling() {
                         break;
@@ -145,6 +151,7 @@ fn traverse_tree(
                 class_name.clone(),
                 module_name,
                 imports,
+                &config,
             );
             if !cursor.goto_next_sibling() {
                 break;
@@ -225,36 +232,144 @@ fn is_import_statement(kind: &str, language: Language) -> bool {
         lang if lang == unsafe { tree_sitter_python() } => {
             kind == "import_statement" || kind == "import_from_statement"
         }
+        lang if lang == unsafe { tree_sitter_rust() } => kind == "use_declaration",
         // Add more language-specific checks here
         _ => false,
     }
 }
 
-fn parse_import_statement(code: &str, node: Node, language: Language) -> Option<(String, String)> {
+fn filter_import_matchers(
+    child: Node,
+    code: &str,
+    matchers: &Matchers,
+) -> (Option<String>, Option<String>, Option<String>) {
+    let module = child
+        .child_by_field_name(&matchers.module_name.field_name)
+        .map(|n| {
+            if n.kind() == matchers.module_name.kind {
+                return n.utf8_text(code.as_bytes()).unwrap_or_default().to_owned();
+            }
+
+            String::default()
+        });
+
+    let name = child
+        .child_by_field_name(&matchers.object_name.field_name)
+        .map(|n| {
+            if n.kind() == matchers.object_name.kind {
+                return n.utf8_text(code.as_bytes()).unwrap_or_default().to_owned();
+            }
+
+            String::default()
+        });
+
+    let alias = child
+        .child_by_field_name(&matchers.alias.field_name)
+        .map(|n| {
+            if n.kind() == matchers.alias.kind {
+                return n.utf8_text(code.as_bytes()).unwrap_or_default().to_owned();
+            }
+
+            String::default()
+        });
+
+    (module, name, alias)
+}
+
+fn parse_import_statement(
+    code: &str,
+    node: Node,
+    language: Language,
+    config: &Config,
+) -> Option<(String, String)> {
+    let mut module_name = String::new();
+    let mut object_name = String::new();
+    let mut alias_name = String::new();
+
     match language {
         lang if lang == unsafe { tree_sitter_python() } => {
-            if node.kind() == "import_from_statement" {
-                // Extract the module name from the 'module_name' field
-                let module = node
-                    .child_by_field_name("module_name")
-                    .map(|n| n.utf8_text(code.as_bytes()).ok())
-                    .flatten()
-                    .unwrap_or_default()
-                    .to_string();
-
-                // Iterate over the names imported from the module
-                let node_walk = &mut node.walk();
-                let imported_names = node.children_by_field_name("name", node_walk);
-                for imported_name in imported_names {
-                    let alias = imported_name
-                        .utf8_text(code.as_bytes())
-                        .unwrap_or_default()
-                        .to_string();
-
-                    // In this case, we assume that each import statement imports a single name,
-                    // so we return the first found. For handling multiple imports, this approach needs to be adjusted.
-                    return Some((module.clone(), alias));
+            let matchers = &config
+            .languages
+            .get("python")
+            .expect("Failed to get Python matchers from config")
+            .matchers;
+
+            if node.kind() == matchers.import_statement {
+                let result = filter_import_matchers(node, code, matchers);
+                (module_name, object_name, alias_name) = (
+                    result.0.unwrap_or(module_name),
+                    result.1.unwrap_or(object_name),
+                    result.2.unwrap_or(alias_name),
+                );
+
+                let mut cursor = node.walk();
+                for child in node.named_children(&mut cursor) {
+                    let result = filter_import_matchers(child, code, matchers);
+                    (module_name, object_name, alias_name) = (
+                        result.0.unwrap_or(module_name),
+                        result.1.unwrap_or(object_name),
+                        result.2.unwrap_or(alias_name),
+                    );
+
+                    let mut cursor2 = child.walk();
+                    for child2 in child.named_children(&mut cursor2) {
+                        let result = filter_import_matchers(child2, code, matchers);
+                        (module_name, object_name, alias_name) = (
+                            result.0.unwrap_or(module_name),
+                            result.1.unwrap_or(object_name),
+                            result.2.unwrap_or(alias_name),
+                        );
+                    }
+                }
+
+                println!(
+                    "Module: {}, Object: {}, Alias: {}",
+                    module_name, object_name, alias_name
+                );
+                return Some((module_name, object_name));
+            }
+            None
+        },
+        lang if lang == unsafe { tree_sitter_rust() } => {
+            let matchers = &config
+            .languages
+            .get("rust")
+            .expect("Failed to get Python matchers from config")
+            .matchers;
+
+            if node.kind() == matchers.import_statement {
+                let result = filter_import_matchers(node, code, matchers);
+                (module_name, object_name, alias_name) = (
+                    result.0.unwrap_or(module_name),
+                    result.1.unwrap_or(object_name),
+                    result.2.unwrap_or(alias_name),
+                );
+
+                let mut cursor = node.walk();
+                for child in node.named_children(&mut cursor) {
+                    let result = filter_import_matchers(child, code, matchers);
+                    (module_name, object_name, alias_name) = (
+                        result.0.unwrap_or(module_name),
+                        result.1.unwrap_or(object_name),
+                        result.2.unwrap_or(alias_name),
+                    );
+
+                    let mut cursor2 = child.walk();
+                    for child2 in child.named_children(&mut cursor2) {
+                        let result = filter_import_matchers(child2, code, matchers);
+                        (module_name, object_name, alias_name) = (
+                            result.0.unwrap_or(module_name),
+                            result.1.unwrap_or(object_name),
+                            result.2.unwrap_or(alias_name),
+                        );
+                    }
                 }
+
+                println!(
+                    "Module: {}, Object: {}, Alias: {}",
+                    module_name, object_name, alias_name
+                );
+                return Some((module_name, object_name));
             }
             None
         }

test-code-base-python/main.py

@@ -1,4 +1,5 @@
 from calc1 import Calc as FakeCalc
+from a import b
 
 def main():
     c = FakeCalc()

test-code-base-python/main2.py

@@ -1,8 +1,10 @@
 mod person;
 mod animal;
 
+use person::Person as Animal;
+
 fn main() {
-    let mut person = person::Person::new("Alice".to_string(), 30);
+    let mut person = Person::new("Alice".to_string(), 30);
     person.introduce();
     person.celebrate_birthday();