fix: python imports

languages/tree-sitter-stack-graphs-python/src/stack-graphs.tsg

@@ -13,6 +13,7 @@
 ;; ^^^^^^^^^^^^^^^^
 
 global FILE_PATH
+global ROOT_PATH = ""
 global ROOT_NODE
 global JUMP_TO_SCOPE_NODE
 
@@ -272,7 +273,9 @@ inherit .parent_module
   node grandparent_module_ref_node
   var grandparent_module_ref = grandparent_module_ref_node
 
-  scan FILE_PATH {
+  ; get the file path relative to the root path
+  let rel_path = (replace FILE_PATH ROOT_PATH "")
+  scan rel_path {
     "([^/]+)/"
     {
       node def_dot

languages/tree-sitter-stack-graphs-python/test/root_path.py

@@ -0,0 +1,21 @@
+# ------ path: foo/bar/module.py -----------#
+# ------ global: ROOT_PATH=foo/bar  -----------#
+
+foo = 42
+
+# ------ path: foo/bar/baz/module.py -----------#
+# ------ global: ROOT_PATH=foo/bar  -----------#
+
+bar = "hello"
+
+# ------ path: foo/bar/main.py -------------#
+# ------ global: ROOT_PATH=foo/bar  -----------#
+
+from module import foo
+from baz.module import bar
+
+print(foo)
+#     ^ defined: 4, 14
+
+print(bar)
+#     ^ defined: 9, 15

tree-sitter-stack-graphs/src/cli/index.rs

@@ -42,6 +42,7 @@ use crate::BuildError;
 use crate::CancelAfterDuration;
 use crate::CancellationFlag;
 use crate::NoCancellation;
+use crate::{FILE_PATH_VAR, ROOT_PATH_VAR};
 
 #[derive(Args)]
 pub struct IndexArgs {
@@ -429,7 +430,16 @@ impl<'a> Indexer<'a> {
     ) -> std::result::Result<(), BuildErrorWithSource<'b>> {
         let relative_source_path = source_path.strip_prefix(source_root).unwrap();
         if let Some(lc) = lcs.primary {
-            let globals = Variables::new();
+            let mut globals = Variables::new();
+
+            globals
+                .add(FILE_PATH_VAR.into(), source_path.to_str().unwrap().into())
+                .expect("failed to add file path variable");
+
+            globals
+                .add(ROOT_PATH_VAR.into(), source_root.to_str().unwrap().into())
+                .expect("failed to add root path variable");
+
             lc.sgl
                 .build_stack_graph_into(graph, file, source, &globals, cancellation_flag)
                 .map_err(|inner| BuildErrorWithSource {

tree-sitter-stack-graphs/src/cli/test.rs

@@ -43,6 +43,7 @@ use crate::test::Test;
 use crate::test::TestResult;
 use crate::CancelAfterDuration;
 use crate::CancellationFlag;
+use crate::FILE_PATH_VAR;
 
 #[derive(Args)]
 #[clap(after_help = r#"PATH SPECIFICATIONS:
@@ -316,6 +317,15 @@ impl TestArgs {
                 &mut Some(test_fragment.source.as_ref()),
             )? {
                 globals.clear();
+
+                // Add FILE_PATH to variables
+                globals
+                    .add(
+                        FILE_PATH_VAR.into(),
+                        test_fragment.path.to_str().unwrap().into(),
+                    )
+                    .expect("failed to add file path variable");
+
                 test_fragment.add_globals_to(&mut globals);
                 lc.sgl.build_stack_graph_into(
                     &mut test.graph,

tree-sitter-stack-graphs/src/lib.rs

@@ -438,6 +438,7 @@ static PRECEDENCE_ATTR: &'static str = "precedence";
 static ROOT_NODE_VAR: &'static str = "ROOT_NODE";
 static JUMP_TO_SCOPE_NODE_VAR: &'static str = "JUMP_TO_SCOPE_NODE";
 static FILE_PATH_VAR: &'static str = "FILE_PATH";
+static ROOT_PATH_VAR: &'static str = "ROOT_PATH";
 
 /// Holds information about how to construct stack graphs for a particular language.
 pub struct StackGraphLanguage {
@@ -635,22 +636,23 @@ impl<'a> Builder<'a> {
         let tree = parse_errors.into_tree();
 
         let mut globals = Variables::nested(globals);
+
         if globals.get(&ROOT_NODE_VAR.into()).is_none() {
             let root_node = self.inject_node(NodeID::root());
             globals
                 .add(ROOT_NODE_VAR.into(), root_node.into())
                 .expect("Failed to set ROOT_NODE");
         }
+
         let jump_to_scope_node = self.inject_node(NodeID::jump_to());
         globals
             .add(JUMP_TO_SCOPE_NODE_VAR.into(), jump_to_scope_node.into())
             .expect("Failed to set JUMP_TO_SCOPE_NODE");
-        if globals.get(&FILE_PATH_VAR.into()).is_none() {
-            let file_name = self.stack_graph[self.file].to_string();
-            globals
-                .add(FILE_PATH_VAR.into(), file_name.into())
-                .expect("Failed to set FILE_PATH");
-        }
+
+        // FILE_PATH is mandatory
+        globals
+            .get(&FILE_PATH_VAR.into())
+            .expect("FILE_PATH not set");
 
         let mut config = ExecutionConfig::new(&self.sgl.functions, &globals)
             .lazy(true)

tree-sitter-stack-graphs/src/loader.rs

@@ -29,6 +29,7 @@ use tree_sitter_loader::Loader as TsLoader;
 use crate::CancellationFlag;
 use crate::FileAnalyzer;
 use crate::StackGraphLanguage;
+use crate::FILE_PATH_VAR;
 
 pub static DEFAULT_TSG_PATHS: Lazy<Vec<LoadPath>> =
     Lazy::new(|| vec![LoadPath::Grammar("queries/stack-graphs".into())]);
@@ -75,6 +76,15 @@ impl LanguageConfiguration {
         let mut builtins = StackGraph::new();
         if let Some((builtins_path, builtins_source)) = builtins_source {
             let mut builtins_globals = Variables::new();
+
+            let builtins_path_var = Path::new("<builtins>");
+            builtins_globals
+                .add(
+                    FILE_PATH_VAR.into(),
+                    builtins_path_var.to_str().unwrap().into(),
+                )
+                .expect("failed to add file path variable");
+
             if let Some(builtins_config) = builtins_config {
                 Loader::load_globals_from_config_str(builtins_config, &mut builtins_globals)?;
             }
@@ -325,8 +335,15 @@ impl Loader {
         graph: &mut StackGraph,
         cancellation_flag: &dyn CancellationFlag,
     ) -> Result<(), LoadError<'a>> {
-        let file = graph.add_file(&path.to_string_lossy()).unwrap();
+        let file_name = path.to_string_lossy();
+        let file: stack_graphs::arena::Handle<stack_graphs::graph::File> =
+            graph.add_file(&file_name).unwrap();
         let mut globals = Variables::new();
+
+        globals
+            .add(FILE_PATH_VAR.into(), path.to_str().unwrap().into())
+            .expect("failed to add file path variable");
+
         Self::load_globals_from_config_str(&config, &mut globals)?;
         sgl.build_stack_graph_into(graph, file, &source, &globals, cancellation_flag)
             .map_err(|err| LoadError::Builtins {

tree-sitter-stack-graphs/tests/it/builder.rs

@@ -5,13 +5,16 @@
 // Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
 // ------------------------------------------------------------------------------------------------
 
+use std::path::Path;
+
 use stack_graphs::graph::StackGraph;
 use tree_sitter_graph::Variables;
 use tree_sitter_stack_graphs::NoCancellation;
 use tree_sitter_stack_graphs::StackGraphLanguage;
 
 use crate::edges::check_stack_graph_edges;
 use crate::nodes::check_stack_graph_nodes;
+use crate::FILE_PATH_VAR;
 
 #[test]
 fn can_support_preexisting_nodes() {
@@ -22,12 +25,18 @@ fn can_support_preexisting_nodes() {
     "#;
     let python = "pass";
 
+    let file_name = "test.py";
+
     let mut graph = StackGraph::new();
-    let file = graph.get_or_create_file("test.py");
+    let file = graph.get_or_create_file(file_name);
     let node_id = graph.new_node_id(file);
     let _preexisting_node = graph.add_scope_node(node_id, true).unwrap();
 
-    let globals = Variables::new();
+    let mut globals = Variables::new();
+    globals
+        .add(FILE_PATH_VAR.into(), file_name.into())
+        .expect("failed to add file path variable");
+
     let language = StackGraphLanguage::from_str(tree_sitter_python::language(), tsg).unwrap();
     language
         .build_stack_graph_into(&mut graph, file, python, &globals, &NoCancellation)
@@ -45,15 +54,21 @@ fn can_support_injected_nodes() {
     "#;
     let python = "pass";
 
+    let file_name = "test.py";
+
     let mut graph = StackGraph::new();
-    let file = graph.get_or_create_file("test.py");
+    let file = graph.get_or_create_file(file_name);
     let node_id = graph.new_node_id(file);
     let _preexisting_node = graph.add_scope_node(node_id, true).unwrap();
 
     let language = StackGraphLanguage::from_str(tree_sitter_python::language(), tsg).unwrap();
     let mut builder = language.builder_into_stack_graph(&mut graph, file, python);
 
     let mut globals = Variables::new();
+    globals
+        .add(FILE_PATH_VAR.into(), file_name.into())
+        .expect("failed to add file path variable");
+
     globals
         .add("EXT_NODE".into(), builder.inject_node(node_id).into())
         .expect("Failed to add EXT_NODE variable");

tree-sitter-stack-graphs/tests/it/main.rs

@@ -5,6 +5,8 @@
 // Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
 // ------------------------------------------------------------------------------------------------
 
+use std::path::Path;
+
 use stack_graphs::arena::Handle;
 use stack_graphs::graph::File;
 use stack_graphs::graph::StackGraph;
@@ -13,6 +15,8 @@ use tree_sitter_stack_graphs::BuildError;
 use tree_sitter_stack_graphs::NoCancellation;
 use tree_sitter_stack_graphs::StackGraphLanguage;
 
+static FILE_PATH_VAR: &'static str = "FILE_PATH";
+
 mod builder;
 mod edges;
 mod loader;
@@ -23,11 +27,19 @@ pub(self) fn build_stack_graph(
     python_source: &str,
     tsg_source: &str,
 ) -> Result<(StackGraph, Handle<File>), BuildError> {
+    let file_name = "test.py";
     let language =
         StackGraphLanguage::from_str(tree_sitter_python::language(), tsg_source).unwrap();
     let mut graph = StackGraph::new();
-    let file = graph.get_or_create_file("test.py");
-    let globals = Variables::new();
+    let file = graph.get_or_create_file(file_name);
+    let mut globals = Variables::new();
+    let source_path = Path::new(file_name);
+
+    // The FILE_PATH_VAR is not accessible here
+    globals
+        .add(FILE_PATH_VAR.into(), source_path.to_str().unwrap().into())
+        .expect("failed to add file path variable");
+
     language.build_stack_graph_into(&mut graph, file, python_source, &globals, &NoCancellation)?;
     Ok((graph, file))
 }

tree-sitter-stack-graphs/tests/it/test.rs

@@ -5,6 +5,7 @@
 // Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
 // ------------------------------------------------------------------------------------------------
 
+use crate::FILE_PATH_VAR;
 use once_cell::sync::Lazy;
 use pretty_assertions::assert_eq;
 use stack_graphs::arena::Handle;
@@ -94,9 +95,18 @@ fn check_test(
         expected_successes + expected_failures,
         assertion_count,
     );
+
     let mut globals = Variables::new();
     for fragments in &test.fragments {
         globals.clear();
+
+        globals
+            .add(
+                FILE_PATH_VAR.into(),
+                fragments.path.to_str().unwrap().into(),
+            )
+            .expect("failed to add file path variable");
+
         fragments.add_globals_to(&mut globals);
         build_stack_graph_into(
             &mut test.graph,