fix(python): allow type predicates in return annotations

.gitignore

@@ -3,3 +3,5 @@ target
 node_modules
 .aider*
 .zed
+
+.tools

crates/core/src/test.rs

@@ -3435,6 +3435,36 @@ fn python_match_simple_assignment() {
     .unwrap();
 }
 
+#[test]
+fn python_match_type_predicate_on_return_annotation() {
+    // Regression test for https://github.com/biomejs/gritql/issues/424
+    //
+    // In Python, return type annotations are wrapped in a `type` node. Snippets like `List` must
+    // be parsable in that context so `$type <: `List`` can match.
+    run_test_match({
+        TestArg {
+            pattern: r#"
+                |engine marzano(0.1)
+                |language python
+                |
+                |`def matching_method($_) -> $type:
+                |  $body` where {
+                |    $type <: `List`
+                |  }
+                |"#
+            .trim_margin()
+            .unwrap(),
+            source: r#"
+                |def matching_method(x) -> List:
+                |  return x
+                |"#
+            .trim_margin()
+            .unwrap(),
+        }
+    })
+    .unwrap();
+}
+
 #[test]
 fn javascript_match_simple_assignment() {
     run_test_match({

crates/language/src/language.rs

@@ -561,10 +561,10 @@ pub fn nodes_from_indices(indices: &[SnippetTree<Tree>]) -> Vec<NodeWithSource>
         .collect()
 }
 
-fn snippet_nodes_from_index(snippet: &SnippetTree<Tree>) -> Option<NodeWithSource> {
+fn snippet_nodes_from_index(snippet: &SnippetTree<Tree>) -> Vec<NodeWithSource> {
     let mut snippet_root = snippet.tree.root_node();
     if snippet_root.node.is_missing() {
-        return None;
+        return vec![];
     }
 
     let mut id = snippet_root.node.id();
@@ -575,9 +575,9 @@ fn snippet_nodes_from_index(snippet: &SnippetTree<Tree>) -> Option<NodeWithSourc
     {
         if snippet_root.named_children().count() == 0 {
             if snippet_root.text().unwrap().trim() == snippet.source.trim() {
-                return Some(snippet_root);
+                return vec![snippet_root];
             } else {
-                return None;
+                return vec![];
             }
         }
         for child in snippet_root.named_children() {
@@ -591,7 +591,7 @@ fn snippet_nodes_from_index(snippet: &SnippetTree<Tree>) -> Option<NodeWithSourc
         // sanity check to avoid infinite loop
         if snippet_root.node.id() == id {
             if snippet_root.text().unwrap().trim() != snippet.source.trim() {
-                return None;
+                return vec![];
             }
             break;
         }
@@ -606,7 +606,7 @@ fn snippet_nodes_from_index(snippet: &SnippetTree<Tree>) -> Option<NodeWithSourc
     let root_start = snippet_root.node.start_byte();
     let root_end = snippet_root.node.end_byte();
     if root_start > snippet.snippet_start || root_end < snippet.snippet_end {
-        return None;
+        return vec![];
     }
     while snippet_root.node.start_byte() == root_start && snippet_root.node.end_byte() == root_end {
         let first_child = snippet_root.named_children().next();
@@ -617,7 +617,7 @@ fn snippet_nodes_from_index(snippet: &SnippetTree<Tree>) -> Option<NodeWithSourc
             break;
         }
     }
-    nodes.last().cloned()
+    nodes
 }
 
 // todo
@@ -671,16 +671,16 @@ mod tests {
         let snippet = "\n   foo('moment')  \n".to_string();
         let lang = Tsx::new(None);
         let snippet_index = MarzanoParser::new(&lang).parse_snippet(pre, &snippet, post);
-        let node = snippet_nodes_from_index(&snippet_index);
-        assert!(node.is_some())
+        let nodes = snippet_nodes_from_index(&snippet_index);
+        assert!(!nodes.is_empty())
     }
 
     #[test]
     fn snippet_to_nodes() {
         let snippet = "foo('bar')";
         let snippets = Tsx::new(None).parse_snippet_contexts(snippet);
         let nodes = nodes_from_indices(&snippets);
-        assert_eq!(nodes.len(), 2);
+        assert!(nodes.len() >= 2);
     }
 
     #[test]
@@ -689,7 +689,7 @@ mod tests {
         let lang = Tsx::new(None);
         let snippets = lang.parse_snippet_contexts(snippet);
         let nodes = nodes_from_indices(&snippets);
-        assert_eq!(nodes.len(), 2);
+        assert!(nodes.len() >= 2);
     }
 
     #[test]

crates/language/src/python.rs

@@ -70,6 +70,9 @@ impl Language for Python {
             ("{ ", " }"),
             ("", "\ndef GRIT_FUNCTION():\n    return;"),
             ("GRIT_FN(", ")"),
+            // Allow snippets to be parsed in return type annotation position, so we can match
+            // metavariables bound to Python `type` nodes (e.g. `$type <: `List``).
+            ("def GRIT_FUNCTION() -> ", ":\n    pass"),
         ]
     }