refactor: simplify and correct Remove op

crates/yamlpath/src/lib.rs

@@ -579,12 +579,15 @@ impl Document {
     fn query_node(&self, query: &Query, mode: QueryMode) -> Result<Node, QueryError> {
         let mut focus_node = self.top_object()?;
         for component in &query.route {
+            dbg!(&focus_node);
             match self.descend(&focus_node, component) {
                 Ok(next) => focus_node = next,
                 Err(e) => return Err(e),
             }
         }
 
+        dbg!(&focus_node);
+
         focus_node = match mode {
             QueryMode::Pretty => {
                 // If we're in "pretty" mode, we want to return the
@@ -594,10 +597,13 @@ impl Document {
                 //
                 // NOTE: We might already be on the block/flow pair if we terminated
                 // with an absent value, in which case we don't need to do this cleanup.
-                if matches!(query.route.last(), Some(Component::Key(_)))
+                if (matches!(query.route.last(), Some(Component::Key(_)))
                     && focus_node.kind_id() != self.block_mapping_pair_id
-                    && focus_node.kind_id() != self.flow_pair_id
+                    && focus_node.kind_id() != self.flow_pair_id)
+                    || (matches!(query.route.last(), Some(Component::Index(_)))
+                        && focus_node.kind_id() == self.flow_node_id)
                 {
+                    dbg!("hmmmm");
                     focus_node.parent().unwrap()
                 } else {
                     focus_node
@@ -650,6 +656,7 @@ impl Document {
         if matches!(mode, QueryMode::Pretty)
             && matches!(query.route.last(), Some(Component::Key(_)))
             && focus_node.kind_id() != self.block_mapping_pair_id
+            && focus_node.kind_id() != self.flow_pair_id
         {
             focus_node = focus_node.parent().unwrap()
         }
@@ -850,34 +857,52 @@ baz: quux
         )
     }
 
+    //     #[test]
+    //     fn test_basic() {
+    //         let doc = r#"
+    // foo: bar
+    // baz:
+    //   sub:
+    //     keys:
+    //       abc:
+    //         - 123
+    //         - 456
+    //         - [a, b, c, {d: e}]
+    //         "#;
+
+    //         let doc = Document::new(doc).unwrap();
+    //         let query = Query {
+    //             route: vec![
+    //                 Component::Key("baz"),
+    //                 Component::Key("sub"),
+    //                 Component::Key("keys"),
+    //                 Component::Key("abc"),
+    //                 Component::Index(2),
+    //                 Component::Index(3),
+    //             ],
+    //         };
+
+    //         assert_eq!(
+    //             doc.extract_with_leading_whitespace(&doc.query_pretty(&query).unwrap()),
+    //             "{d: e}"
+    //         );
+    //     }
+
     #[test]
-    fn test_basic() {
+    fn test_flow_pair_in_sequence() {
         let doc = r#"
-foo: bar
-baz:
-  sub:
-    keys:
-      abc:
-        - 123
-        - 456
-        - [a, b, c, {d: e}]
-        "#;
+foo: [1, 2, 3: {abc: def}]
+"#;
 
         let doc = Document::new(doc).unwrap();
         let query = Query {
-            route: vec![
-                Component::Key("baz"),
-                Component::Key("sub"),
-                Component::Key("keys"),
-                Component::Key("abc"),
-                Component::Index(2),
-                Component::Index(3),
-            ],
+            route: vec![Component::Key("foo"), Component::Index(2)],
         };
 
+        let feature = doc.query_pretty(&query).unwrap();
         assert_eq!(
-            doc.extract_with_leading_whitespace(&doc.query_pretty(&query).unwrap()),
-            "{d: e}"
+            doc.extract_with_leading_whitespace(&feature),
+            "3: {abc: def}"
         );
     }
 

crates/yamlpath/tests/testcases/flow.yml

@@ -12,11 +12,11 @@ testcase:
 
 queries:
   - query: [flow1, foo]
-    expected: "{ foo: [1, 2, 3: [4, 5, { a: b }]] }"
+    expected: "foo: [1, 2, 3: [4, 5, { a: b }]]"
 
   - query: [flow1, foo, 2]
     # TODO: ideally would be `3: [4, 5, { a: b }]`
-    expected: "[4, 5, { a: b }]"
+    expected: "3: [4, 5, { a: b }]"
 
   - query: [flow1, foo, 2, 2]
     expected: "{ a: b }"
@@ -28,7 +28,7 @@ queries:
     expected: "{ a: }"
 
   - query: [flow4, foo]
-    expected: "{\n    foo: [1, 2, 3: [4, 5, { a: b }]],\n  }"
+    expected: "    foo: [1, 2, 3: [4, 5, { a: b }]]"
 
   - query: [flow5, 0]
     expected: "    abc"
@@ -37,4 +37,4 @@ queries:
     expected: "def"
 
   - query: [flow5, 2]
-    expected: "    xyz"
+    expected: "    xyz"

crates/zizmor/src/yaml_patch/mod.rs

@@ -168,7 +168,6 @@ pub enum Op<'doc> {
         updates: indexmap::IndexMap<String, serde_yaml::Value>,
     },
     /// Remove the key at the given path
-    #[allow(dead_code)]
     Remove,
 }
 
@@ -436,20 +435,16 @@ fn apply_single_patch(
 
             let feature = route_to_feature_pretty(&patch.route, document)?;
 
-            // For removal, we need to remove the entire line including leading whitespace
-            // TODO: This isn't sound, e.g. removing `b:` from `{a: a, b: b}` will
-            // remove the entire line.
-            let start_pos = {
-                let range = line_span(document, feature.location.byte_span.0);
-                range.start
-            };
-            let end_pos = {
-                let range = line_span(document, feature.location.byte_span.1);
-                range.end
-            };
-
+            // At the moment, we simply delete the entire "pretty" feature.
+            // This works well enough, but it will leave any leading
+            // or trailing whitespace if not captured in the underlying
+            // tree-sitter-extracted feature.
             let mut result = content.to_string();
-            result.replace_range(start_pos..end_pos, "");
+            result.replace_range(
+                feature.location.byte_span.0..feature.location.byte_span.1,
+                "",
+            );
+
             yamlpath::Document::new(result).map_err(Error::from)
         }
     }
@@ -1500,7 +1495,7 @@ foo: { bar: abc }
     }
 
     #[test]
-    fn test_remove_preserves_structure() {
+    fn test_remove_preserves_block_structure() {
         let original = r#"
 permissions:
   contents: read  # Keep this comment
@@ -1525,6 +1520,47 @@ permissions:
         ");
     }
 
+    // #[test]
+    // fn test_remove_preserves_flow_structure() {
+    //     let original = "foo: { a: a, b: b }";
+
+    //     let document = yamlpath::Document::new(original).unwrap();
+
+    //     let operations = vec![Patch {
+    //         route: route!("foo", "b"),
+    //         operation: Op::Remove,
+    //     }];
+
+    //     let result = apply_yaml_patches(&document, &operations).unwrap();
+
+    //     // Removes the key while preserving the rest of the structure
+    //     insta::assert_snapshot!(result.source(), @"foo: { a: a }");
+    // }
+
+    #[test]
+    fn test_remove_empty_key() {
+        let original = r#"
+foo:
+  bar: abc
+  baz:
+"#;
+
+        let document = yamlpath::Document::new(original).unwrap();
+
+        let operations = vec![Patch {
+            route: route!("foo", "baz"),
+            operation: Op::Remove,
+        }];
+
+        let result = apply_yaml_patches(&document, &operations).unwrap();
+
+        // Removes the empty key while preserving the rest of the structure
+        insta::assert_snapshot!(result.source(), @r"
+        foo:
+          bar: abc
+        ");
+    }
+
     #[test]
     fn test_multiple_operations_preserve_comments() {
         let original = r#"