Added line only matching option to XMLTransformer (#733)

pixee · Jul 23, 2024 · 3a20d20 · 3a20d20
1 parent 5f4f575
commit 3a20d20
Show file tree

Hide file tree

Showing 2 changed files with 87 additions and 7 deletions.
diff --git a/src/codemodder/codemods/xml_transformer.py b/src/codemodder/codemods/xml_transformer.py
@@ -30,11 +30,13 @@ def __init__(
         encoding: str = "utf-8",
         short_empty_elements: bool = False,
         results: list[Result] | None = None,
+        line_only_matching=False,
     ) -> None:
         self.file_context = file_context
         self.results = results
         self.changes: list[Change] = []
         self._my_locator = Locator()
+        self.line_only_matching = line_only_matching
         super().__init__(out, encoding, short_empty_elements)
 
     def startElement(self, name, attrs):
@@ -81,10 +83,15 @@ def match_result(self, line, column) -> bool:
             return True
         for result in self.results or []:
             for location in result.locations:
-                # No two elements can have the same start but different ends.
-                # It suffices to only match the start.
-                if location.start.line == line and location.start.column - 1 == column:
-                    return True
+                if self.line_only_matching:
+                    return location.start.line == line
+                else:
+                    # No two elements can have the same start but different ends.
+                    # It suffices to only match the start.
+                    return (
+                        location.start.line == line
+                        and location.start.column - 1 == column
+                    )
         return False
 
     def add_change(self, line):
@@ -110,9 +117,17 @@ def __init__(
         encoding: str = "utf-8",
         short_empty_elements: bool = False,
         results: list[Result] | None = None,
+        line_only_matching=False,
     ) -> None:
         self.name_attributes_map = name_attributes_map
-        super().__init__(out, file_context, encoding, short_empty_elements, results)
+        super().__init__(
+            out,
+            file_context,
+            encoding,
+            short_empty_elements,
+            results,
+            line_only_matching,
+        )
 
     def startElement(self, name, attrs):
         new_attrs: AttributesImpl = attrs

diff --git a/tests/test_xml_transformer.py b/tests/test_xml_transformer.py
@@ -13,6 +13,7 @@
     NewElementXMLTransformer,
     XMLTransformer,
 )
+from codemodder.semgrep import SemgrepResult
 
 
 class TestXMLTransformer:
@@ -56,10 +57,21 @@ def test_parse_cdata(self):
 
 class TestElementAttributeXMLTransformer:
 
-    def run_and_assert(self, name_attr_map, input_code, expected_output):
+    def run_and_assert(
+        self,
+        name_attr_map,
+        input_code,
+        expected_output,
+        results=None,
+        line_only_match=False,
+    ):
         with StringIO() as result, StringIO(dedent(input_code)) as input_stream:
             transformer = ElementAttributeXMLTransformer(
-                result, mock.MagicMock(), name_attributes_map=name_attr_map
+                result,
+                mock.MagicMock(),
+                name_attributes_map=name_attr_map,
+                results=results,
+                line_only_matching=line_only_match,
             )
             parser = make_parser()
             parser.setContentHandler(transformer)
@@ -97,6 +109,59 @@ def test_change_multiple_attr_and_preserve_existing(self):
         name_attr_map = {"element": {"first": "one", "second": "two"}}
         self.run_and_assert(name_attr_map, input_code, expected_output)
 
+    def test_change_only_line_matching_result(self):
+        input_code = """\
+        <?xml version="1.0" encoding="utf-8"?>
+        <configuration>
+          <element first="1">
+          </element>
+          <element first="1">
+          </element>
+        </configuration>"""
+        expected_output = """\
+        <?xml version="1.0" encoding="utf-8"?>
+        <configuration>
+          <element first="1">
+          </element>
+          <element first="one">
+          </element>
+        </configuration>"""
+        name_attr_map = {"element": {"first": "one"}}
+        data = {
+            "runs": [
+                {
+                    "results": [
+                        {
+                            "fingerprints": {"matchBasedId/v1": "123"},
+                            "locations": [
+                                {
+                                    "ruleId": "rule",
+                                    "physicalLocation": {
+                                        "artifactLocation": {
+                                            "uri": "code.py",
+                                            "uriBaseId": "%SRCROOT%",
+                                        },
+                                        "region": {
+                                            "snippet": {"text": "snip"},
+                                            "endColumn": 1,
+                                            "endLine": 5,
+                                            "startColumn": 1,
+                                            "startLine": 5,
+                                        },
+                                    },
+                                }
+                            ],
+                            "ruleId": "rule",
+                        }
+                    ]
+                }
+            ]
+        }
+        sarif_run = data["runs"]
+        sarif_results = sarif_run[0]["results"]
+        results = [SemgrepResult.from_sarif(sarif_results[0], sarif_run)]
+        self.run_and_assert(name_attr_map, input_code, expected_output, results, True)
+
 
 class TestNewElementXMLTransformer: