feat: support more languages

Author: nohehf

Cargo.lock

@@ -12,5 +12,7 @@ crate-type = ["cdylib"]
 pyo3 = "0.20.0"
 stack-graphs = { version = "0.13.0", features = ["storage"] }
 tree-sitter-stack-graphs = { version = "0.8.1", features = ["cli"] }
+tree-sitter-stack-graphs-java = "0.3.0"
 tree-sitter-stack-graphs-javascript = "0.1.0"
 tree-sitter-stack-graphs-python = "0.1.0"
+tree-sitter-stack-graphs-typescript = "0.2.0"

Cargo.toml

@@ -0,0 +1,60 @@
+use std::fmt::Display;
+
+use pyo3::prelude::*;
+use tree_sitter_stack_graphs::cli::util::{SourcePosition, SourceSpan};
+
+#[pyclass]
+#[derive(Clone)]
+pub enum Language {
+    Python,
+    JavaScript,
+    TypeScript,
+    Java,
+}
+
+#[pyclass]
+#[derive(Clone)]
+pub struct Position {
+    #[pyo3(get, set)]
+    path: String,
+    #[pyo3(get, set)]
+    line: usize,
+    #[pyo3(get, set)]
+    column: usize,
+}
+
+#[pymethods]
+impl Position {
+    #[new]
+    fn new(path: String, line: usize, column: usize) -> Self {
+        Position { path, line, column }
+    }
+}
+
+impl Display for Position {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        write!(f, "{}:{}:{}", self.path, self.line, self.column)
+    }
+}
+
+impl From<Position> for SourcePosition {
+    fn from(pos: Position) -> Self {
+        SourcePosition {
+            path: pos.path.into(),
+            line: pos.line,
+            column: pos.column,
+        }
+    }
+}
+
+impl Into<Position> for SourceSpan {
+    fn into(self) -> Position {
+        Position {
+            // TODO(@nohehf): Unwrap is unsafe
+            path: self.path.to_str().unwrap().to_string(),
+            line: self.span.start.line,
+            // TODO(@nohehf): Handle both utf8 and utf16
+            column: self.span.start.column.utf8_offset,
+        }
+    }
+}

src/classes.rs

@@ -1,10 +1,9 @@
-use core::fmt::Display;
-
 use pyo3::prelude::*;
 
+mod classes;
 mod stack_graphs_wrapper;
 
-use tree_sitter_stack_graphs::cli::util::{SourcePosition, SourceSpan};
+use classes::{Language, Position};
 
 /// Formats the sum of two numbers as string.
 #[pyfunction]
@@ -13,61 +12,19 @@ fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
 }
 /// Indexes the given paths into stack graphs, and stores the results in the given database.
 #[pyfunction]
-fn index(paths: Vec<String>, db_path: String) -> PyResult<()> {
+fn index(paths: Vec<String>, db_path: String, language: Language) -> PyResult<()> {
+    // TODO(@nohehf): Add a verbose mode to toggle the logs
     println!("Indexing paths: {:?}", paths);
     println!("Database path: {:?}", db_path);
 
     let paths: Vec<std::path::PathBuf> =
         paths.iter().map(|p| std::path::PathBuf::from(p)).collect();
 
-    Ok(stack_graphs_wrapper::index(paths, &db_path)?)
-}
-
-#[pyclass]
-#[derive(Clone)]
-struct Position {
-    #[pyo3(get, set)]
-    path: String,
-    #[pyo3(get, set)]
-    line: usize,
-    #[pyo3(get, set)]
-    column: usize,
-}
-
-#[pymethods]
-impl Position {
-    #[new]
-    fn new(path: String, line: usize, column: usize) -> Self {
-        Position { path, line, column }
-    }
-}
-
-impl Display for Position {
-    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "{}:{}:{}", self.path, self.line, self.column)
-    }
-}
-
-impl From<Position> for SourcePosition {
-    fn from(pos: Position) -> Self {
-        SourcePosition {
-            path: pos.path.into(),
-            line: pos.line,
-            column: pos.column,
-        }
-    }
-}
-
-impl Into<Position> for SourceSpan {
-    fn into(self) -> Position {
-        Position {
-            // TODO(@nohehf): Unwrap is unsafe
-            path: self.path.to_str().unwrap().to_string(),
-            line: self.span.start.line,
-            // TODO(@nohehf): Handle both utf8 and utf16
-            column: self.span.start.column.utf8_offset,
-        }
-    }
+    Ok(stack_graphs_wrapper::index(
+        paths,
+        &db_path,
+        language.into(),
+    )?)
 }
 
 /// Indexes the given paths into stack graphs, and stores the results in the given database.
@@ -96,5 +53,6 @@ fn stack_graphs_python(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(index, m)?)?;
     m.add_function(wrap_pyfunction!(query_definition, m)?)?;
     m.add_class::<Position>()?;
+    m.add_class::<Language>()?;
     Ok(())
 }

src/lib.rs

@@ -1,10 +1,12 @@
+use crate::Language;
 use pyo3::exceptions::PyException;
 use pyo3::PyErr;
 use stack_graphs::storage::{SQLiteReader, SQLiteWriter};
 use std::path::PathBuf;
 use tree_sitter_stack_graphs::cli::query::{Querier, QueryResult};
 use tree_sitter_stack_graphs::cli::util::SourcePosition;
 use tree_sitter_stack_graphs::cli::{index::Indexer, util::reporter::ConsoleReporter};
+use tree_sitter_stack_graphs::loader::LanguageConfiguration;
 use tree_sitter_stack_graphs::{loader::Loader, NoCancellation};
 
 // TODO(@nohehf): Better error handling
@@ -19,13 +21,29 @@ impl std::convert::From<StackGraphsError> for PyErr {
     }
 }
 
-pub fn index(paths: Vec<PathBuf>, db_path: &str) -> Result<(), StackGraphsError> {
-    let py_config = tree_sitter_stack_graphs_python::language_configuration(&NoCancellation);
-    let js_config = tree_sitter_stack_graphs_javascript::language_configuration(&NoCancellation);
+fn get_langauge_configuration(lang: Language) -> LanguageConfiguration {
+    match lang {
+        Language::Python => {
+            tree_sitter_stack_graphs_python::language_configuration(&NoCancellation)
+        }
+        Language::JavaScript => {
+            tree_sitter_stack_graphs_javascript::language_configuration(&NoCancellation)
+        }
+        Language::TypeScript => {
+            tree_sitter_stack_graphs_typescript::language_configuration(&NoCancellation)
+        }
+        Language::Java => tree_sitter_stack_graphs_java::language_configuration(&NoCancellation),
+    }
+}
 
-    let configs = vec![js_config, py_config];
+pub fn index(
+    paths: Vec<PathBuf>,
+    db_path: &str,
+    language: Language,
+) -> Result<(), StackGraphsError> {
+    let configurations = vec![get_langauge_configuration(language)];
 
-    let mut loader = match Loader::from_language_configurations(configs, None) {
+    let mut loader = match Loader::from_language_configurations(configurations, None) {
         Ok(ldr) => ldr,
         Err(e) => {
             return Err(StackGraphsError {
@@ -81,16 +99,6 @@ pub fn query_definition(
             message: format!("Failed to query definitions: {}", e),
         }),
     }
-
-    // if results.is_empty() {
-    //     println!("No definitions found");
-    //     return Ok(());
-    // }
-
-    // for res in results {
-    //     println!("Source: {:?}", res.source);
-    //     println!("Targets: {:?}", res.targets);
-    // }
 }
 
 // https://github.com/github/stack-graphs/blob/7db914c01b35ce024f6767e02dd1ad97022a6bc1/tree-sitter-stack-graphs/src/cli/index.rs#L118