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

@@ -435,9 +435,16 @@ static SCOPE_ATTRS: Lazy<HashSet<&'static str>> =
 static PRECEDENCE_ATTR: &'static str = "precedence";
 
 // Global variables
-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";
+/// Name of the variable used to pass the root node.
+pub const ROOT_NODE_VAR: &'static str = "ROOT_NODE";
+/// Name of the variable used to pass the jump-to-scope node.
+pub const JUMP_TO_SCOPE_NODE_VAR: &'static str = "JUMP_TO_SCOPE_NODE";
+/// Name of the variable used to pass the file path.
+/// If a root path is given, it should be a descendant the root path.
+pub const FILE_PATH_VAR: &'static str = "FILE_PATH";
+/// Name of the variable used to pass the root path.
+/// If given, should be an ancestor of the file path.
+pub const ROOT_PATH_VAR: &'static str = "ROOT_PATH";
 
 /// Holds information about how to construct stack graphs for a particular language.
 pub struct StackGraphLanguage {
@@ -635,22 +642,18 @@ 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");
-        }
 
         let mut config = ExecutionConfig::new(&self.sgl.functions, &globals)
             .lazy(true)

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

@@ -29,6 +29,9 @@ use tree_sitter_loader::Loader as TsLoader;
 use crate::CancellationFlag;
 use crate::FileAnalyzer;
 use crate::StackGraphLanguage;
+use crate::FILE_PATH_VAR;
+
+const BUILTINS_FILENAME: &str = "<builtins>";
 
 pub static DEFAULT_TSG_PATHS: Lazy<Vec<LoadPath>> =
     Lazy::new(|| vec![LoadPath::Grammar("queries/stack-graphs".into())]);
@@ -75,10 +78,15 @@ impl LanguageConfiguration {
         let mut builtins = StackGraph::new();
         if let Some((builtins_path, builtins_source)) = builtins_source {
             let mut builtins_globals = Variables::new();
+
+            builtins_globals
+                .add(FILE_PATH_VAR.into(), BUILTINS_FILENAME.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)?;
             }
-            let file = builtins.add_file("<builtins>").unwrap();
+            let file = builtins.add_file(BUILTINS_FILENAME).unwrap();
             sgl.build_stack_graph_into(
                 &mut builtins,
                 file,
@@ -325,8 +333,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(), BUILTINS_FILENAME.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

@@ -4,11 +4,11 @@
 // Licensed under either of Apache License, Version 2.0, or MIT license, at your option.
 // Please see the LICENSE-APACHE or LICENSE-MIT files in this distribution for license details.
 // ------------------------------------------------------------------------------------------------
-
 use stack_graphs::graph::StackGraph;
 use tree_sitter_graph::Variables;
 use tree_sitter_stack_graphs::NoCancellation;
 use tree_sitter_stack_graphs::StackGraphLanguage;
+use tree_sitter_stack_graphs::FILE_PATH_VAR;
 
 use crate::edges::check_stack_graph_edges;
 use crate::nodes::check_stack_graph_nodes;
@@ -22,12 +22,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 +51,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,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::arena::Handle;
 use stack_graphs::graph::File;
 use stack_graphs::graph::StackGraph;
 use tree_sitter_graph::Variables;
 use tree_sitter_stack_graphs::BuildError;
 use tree_sitter_stack_graphs::NoCancellation;
 use tree_sitter_stack_graphs::StackGraphLanguage;
+use tree_sitter_stack_graphs::FILE_PATH_VAR;
 
 mod builder;
 mod edges;
@@ -23,11 +26,18 @@ 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);
+
+    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,