feat: index

Cargo.toml

@@ -10,3 +10,7 @@ crate-type = ["cdylib"]
 
 [dependencies]
 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-javascript = "0.1.0"
+tree-sitter-stack-graphs-python = "0.1.0"

src/lib.rs

@@ -1,16 +1,25 @@
 use pyo3::prelude::*;
 
-mod stack_graphs;
+mod stack_graphs_wrapper;
 
 /// Formats the sum of two numbers as string.
 #[pyfunction]
 fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
     Ok((a + b).to_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<()> {
+    let paths: Vec<std::path::PathBuf> =
+        paths.iter().map(|p| std::path::PathBuf::from(p)).collect();
+
+    Ok(stack_graphs_wrapper::index(paths, &db_path)?)
+}
 
 /// A Python module implemented in Rust.
 #[pymodule]
 fn stack_graphs_python(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
+    m.add_function(wrap_pyfunction!(index, m)?)?;
     Ok(())
 }

src/stack_graphs_wrapper/mod.rs

@@ -0,0 +1,106 @@
+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::Loader, NoCancellation};
+
+// TODO(@nohehf): Better error handling
+#[derive(Debug, Clone)]
+pub struct IndexError {
+    message: String,
+}
+
+impl std::convert::From<IndexError> for PyErr {
+    fn from(err: IndexError) -> PyErr {
+        PyException::new_err(err.message)
+    }
+}
+
+pub fn index(paths: Vec<PathBuf>, db_path: &str) -> Result<(), IndexError> {
+    let py_config = tree_sitter_stack_graphs_python::language_configuration(&NoCancellation);
+    let js_config = tree_sitter_stack_graphs_javascript::language_configuration(&NoCancellation);
+
+    let configs = vec![js_config, py_config];
+
+    let mut loader = match Loader::from_language_configurations(configs, None) {
+        Ok(ldr) => ldr,
+        Err(e) => {
+            return Err(IndexError {
+                message: format!("Failed to create loader: {}", e),
+            });
+        }
+    };
+
+    let mut db_write = match SQLiteWriter::open(&db_path) {
+        Ok(db) => db,
+        Err(e) => {
+            return Err(IndexError {
+                message: format!("Failed to open database: {}", e),
+            });
+        }
+    };
+
+    let reporter = ConsoleReporter::none();
+
+    let mut indexer = Indexer::new(&mut db_write, &mut loader, &reporter);
+
+    // For now, force reindexing
+    indexer.force = true;
+
+    // TODO(@nohehf): Pass this as input
+    let source_paths: Vec<PathBuf> = vec!["/Users/nohehf/tmp/js".into()];
+
+    let source_paths = canonicalize_paths(source_paths);
+
+    // https://github.com/github/stack-graphs/blob/7db914c01b35ce024f6767e02dd1ad97022a6bc1/tree-sitter-stack-graphs/src/cli/index.rs#L107
+    let continue_from_none: Option<PathBuf> = None;
+
+    match indexer.index_all(source_paths, continue_from_none, &NoCancellation) {
+        Ok(_) => Ok(()),
+        Err(e) => Err(IndexError {
+            message: format!("Failed to index: {}", e),
+        }),
+    }
+}
+
+pub fn query_definition(
+    reference: SourcePosition,
+    db_path: &str,
+) -> Result<Vec<QueryResult>, IndexError> {
+    let mut db_read = SQLiteReader::open(&db_path).expect("failed to open database");
+
+    let reporter = ConsoleReporter::none();
+
+    let mut querier = Querier::new(&mut db_read, &reporter);
+
+    // print_source_position(&reference);
+
+    match querier.definitions(reference, &NoCancellation) {
+        Ok(results) => Ok(results),
+        Err(e) => Err(IndexError {
+            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
+fn canonicalize_paths(paths: Vec<PathBuf>) -> Vec<PathBuf> {
+    paths
+        .into_iter()
+        .map(|p| p.canonicalize())
+        .collect::<std::result::Result<Vec<_>, _>>()
+        .unwrap()
+}

stack_graphs_python.pyi

@@ -1 +1 @@
-def sum_as_string(a: int, b: int) -> str: ...
+def index(paths: list[str], db_path: str) -> None: ...

tests/js_sample/index.js

@@ -0,0 +1,5 @@
+import { foo } from "./module"
+
+const baz = foo
+
+console.log(baz)

tests/js_sample/module.js

@@ -0,0 +1 @@
+export const foo = "bar"

tests/test.py

@@ -1,3 +1,9 @@
-from stack_graphs_python import sum_as_string
+import os
+from stack_graphs_python import index
 
-print(sum_as_string(1, 2))  # '3'
+# index ./js_sample directory
+
+# convert ./js_sample directory to absolute path
+dir = os.path.abspath("./js_sample")
+
+index([dir], "./js_sample.db")