From e12de3b59978e462269263088ef834b90e0be244 Mon Sep 17 00:00:00 2001
From: Alexander Ley <94057608+Alex-ley-scrub@users.noreply.github.com>
Date: Sat, 26 Oct 2024 23:48:26 +0100
Subject: [PATCH] fix: 416 python hanging comma (#566)

Co-authored-by: Morgante Pell <morgantep@google.com>
---
 crates/core/src/test.rs       | 60 +++++++++++++++++++++++++++++++++++
 crates/language/src/python.rs | 13 +++++++-
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/crates/core/src/test.rs b/crates/core/src/test.rs
index d4976a848..640fbef9c 100644
--- a/crates/core/src/test.rs
+++ b/crates/core/src/test.rs
@@ -12272,6 +12272,66 @@ fn trailing_comma_import_from_python_with_alias() {
     .unwrap();
 }
 
+// refer to https://github.com/getgrit/gritql/issues/416
+#[test]
+fn trailing_comma_after_argument_removal() {
+    run_test_expected({
+        TestArgExpected {
+            pattern: r#"
+                language python
+                `TaskMetadata($args)` where {
+                        $args <: contains `n_samples=$_` => .
+                    }
+                "#.to_owned(),
+            source: r#"
+                TaskMetadata(
+                    description="Parallel news titles from the Tbilisi City Hall website (https://tbilisi.gov.ge/).",
+                    main_score="f1", domains=["News"],
+                    text_creation="created",
+                    n_samples={_EVAL_SPLIT: 1820},
+                    reference="https://huggingface.co/datasets/jupyterjazz/tbilisi-city-hall-titles"
+                )
+            "#
+                .to_owned(),
+            expected: r#"
+                TaskMetadata(
+                    description="Parallel news titles from the Tbilisi City Hall website (https://tbilisi.gov.ge/).",
+                    main_score="f1", domains=["News"],
+                    text_creation="created",
+
+                    reference="https://huggingface.co/datasets/jupyterjazz/tbilisi-city-hall-titles"
+                )
+            "#
+                .to_owned(),
+        }
+    })
+    .unwrap();
+}
+
+/// Same as above test, but ensures the behavior doesn't depend on line breaks
+#[test]
+fn trailing_comma_after_argument_removal_one_line() {
+    run_test_expected({
+        TestArgExpected {
+            pattern: r#"
+                language python
+                `TaskMetadata($args)` where {
+                        $args <: contains `n_samples=$_` => .
+                    }
+                "#.to_owned(),
+            source: r#"
+                TaskMetadata(description="Parallel news titles from the Tbilisi City Hall website (https://tbilisi.gov.ge/).", main_score="f1", domains=["News"], text_creation="created", n_samples={_EVAL_SPLIT: 1820}, reference="https://huggingface.co/datasets/jupyterjazz/tbilisi-city-hall-titles")
+            "#
+                .to_owned(),
+            expected: r#"
+                TaskMetadata(description="Parallel news titles from the Tbilisi City Hall website (https://tbilisi.gov.ge/).", main_score="f1", domains=["News"], text_creation="created",  reference="https://huggingface.co/datasets/jupyterjazz/tbilisi-city-hall-titles")
+            "#
+                .to_owned(),
+        }
+    })
+    .unwrap();
+}
+
 #[test]
 fn python_orphaned_from_imports() {
     run_test_expected({
diff --git a/crates/language/src/python.rs b/crates/language/src/python.rs
index 52efd2d80..5e0171f8a 100644
--- a/crates/language/src/python.rs
+++ b/crates/language/src/python.rs
@@ -79,11 +79,22 @@ impl Language for Python {
 
     fn check_replacements(&self, n: NodeWithSource<'_>, replacements: &mut Vec<Replacement>) {
         if n.node.is_error() {
-            if n.text().is_ok_and(|t| t == "->") {
+            // https://github.com/getgrit/gritql/issues/416 single line example
+            if n.text().is_ok_and(|t| t == "->") || n.text().is_ok_and(|t| t == ",") {
                 replacements.push(Replacement::new(n.range(), ""));
             }
             return;
         }
+        if n.node.kind() == "," {
+            // https://github.com/getgrit/gritql/issues/416 multi line example
+            if let Some(prev) = n.node.prev_sibling() {
+                let empty = prev.range().end_byte() == prev.range().start_byte();
+                if empty {
+                    replacements.push(Replacement::new(n.range(), ""));
+                    return;
+                }
+            }
+        }
         if n.node.kind() == "import_from_statement" {
             if let Some(name_field) = n.node.child_by_field_name("name") {
                 let names_text = name_field