fix: remove orphaned type annotation (#122)

getgrit · Mar 29, 2024 · 1fe2884 · 1fe2884
1 parent 0ca9be8
commit 1fe2884
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 0 deletions.
diff --git a/crates/core/src/test.rs b/crates/core/src/test.rs
@@ -7700,6 +7700,46 @@ fn multiply_decimals() {
     .unwrap();
 }
 
+#[test]
+fn python_removes_orphaned_type_arrow() {
+    run_test_expected({
+        TestArgExpected {
+            pattern: r#"
+                |language python
+                |
+                |function_definition($name, $return_type) where {
+                |  $name <: `foo`,
+                |  $return_type => .
+                |}
+                |"#
+            .trim_margin()
+            .unwrap(),
+            source: r#"
+                |def foo() -> None:
+                |    print('hi')
+                |
+                |def bar() -> SomeType:
+                |    print('hello')
+                |    return SomeType
+                |"#
+            .trim_margin()
+            .unwrap(),
+            // The whitespace is fine, because Ruff will remove it
+            expected: r#"
+                |def foo()  :
+                |    print('hi')
+                |
+                |def bar() -> SomeType:
+                |    print('hello')
+                |    return SomeType
+                |"#
+            .trim_margin()
+            .unwrap(),
+        }
+    })
+    .unwrap();
+}
+
 #[test]
 fn removes_orphaned_semicolon() {
     run_test_expected({

diff --git a/crates/language/src/python.rs b/crates/language/src/python.rs
@@ -1,5 +1,7 @@
 use std::sync::OnceLock;
 
+use tree_sitter::Node;
+
 use crate::language::{fields_for_nodes, Field, Language, SortId, TSLanguage};
 
 static NODE_TYPES_STRING: &str =
@@ -77,6 +79,17 @@ impl Language for Python {
     fn is_comment(&self, id: SortId) -> bool {
         id == self.comment_sort
     }
+
+    fn check_orphaned(&self, n: Node<'_>, src: &str, orphan_ranges: &mut Vec<tree_sitter::Range>) {
+        if n.is_error() {
+            let Ok(text) = n.utf8_text(src.as_bytes()) else {
+                return;
+            };
+            if &text == "->" {
+                orphan_ranges.push(n.range());
+            }
+        }
+    }
 }
 
 #[cfg(test)]