Xml transformer (#598)

* xml transformer can add new element

* can add nested elements

* cleanup

Author: clavedeluna

src/codemodder/codemods/xml_transformer.py

@@ -1,4 +1,5 @@
 import mmap
+from dataclasses import dataclass, field
 from tempfile import TemporaryFile
 from xml.sax import SAXParseException, handler
 from xml.sax.handler import LexicalHandler
@@ -114,6 +115,50 @@ def startElement(self, name, attrs):
         super().startElement(name, new_attrs)
 
 
+@dataclass
+class NewElement:
+    name: str
+    parent_name: str
+    content: str = ""
+    attributes: dict[str, str] = field(default_factory=dict)
+
+
+class NewElementXMLTransformer(XMLTransformer):
+    """
+    Adds new elements to the XML file at specified locations.
+    """
+
+    def __init__(
+        self,
+        out,
+        encoding: str = "utf-8",
+        short_empty_elements: bool = False,
+        results: list[Result] | None = None,
+        new_elements: list[NewElement] | None = None,
+    ) -> None:
+        super().__init__(out, encoding, short_empty_elements, results)
+        self.new_elements = new_elements or []
+
+    def startElement(self, name, attrs):
+        super().startElement(name, attrs)
+
+    def endElement(self, name):
+        for new_element in self.new_elements:
+            if new_element.parent_name == name:
+                self.add_new_element(new_element)
+                self.add_change(self._my_locator.getLineNumber())
+        super().endElement(name)
+
+    def add_new_element(self, new_element: NewElement):
+        attrs = AttributesImpl(new_element.attributes or {})
+        super().startElement(new_element.name, attrs)
+        if isinstance(new_element.content, NewElement):
+            self.add_new_element(new_element.content)
+        else:
+            super().characters(new_element.content)
+        super().endElement(new_element.name)
+
+
 class XMLTransformerPipeline(BaseTransformerPipeline):
 
     def __init__(self, xml_transformer: type[XMLTransformer]):

tests/test_xml_transformer.py

@@ -8,6 +8,8 @@
 
 from codemodder.codemods.xml_transformer import (
     ElementAttributeXMLTransformer,
+    NewElement,
+    NewElementXMLTransformer,
     XMLTransformer,
 )
 
@@ -55,7 +57,6 @@ class TestElementAttributeXMLTransformer:
 
     def run_and_assert(self, name_attr_map, input_code, expected_output):
         with StringIO() as result, StringIO(dedent(input_code)) as input_stream:
-            result = StringIO()
             transformer = ElementAttributeXMLTransformer(
                 result, name_attributes_map=name_attr_map
             )
@@ -94,3 +95,91 @@ def test_change_multiple_attr_and_preserve_existing(self):
                 <element first="one" second="two" three="three"></element>"""
         name_attr_map = {"element": {"first": "one", "second": "two"}}
         self.run_and_assert(name_attr_map, input_code, expected_output)
+
+
+class TestNewElementXMLTransformer:
+
+    def run_and_assert(self, new_elements, input_code, expected_output):
+        with StringIO() as result, StringIO(dedent(input_code)) as input_stream:
+            transformer = NewElementXMLTransformer(result, new_elements=new_elements)
+            parser = make_parser()
+            parser.setContentHandler(transformer)
+            parser.setProperty(handler.property_lexical_handler, transformer)
+            parser.parse(input_stream)
+            assert result.getvalue() == dedent(expected_output)
+
+    def test_add_new_element(self):
+        input_code = """\
+                <root></root>
+        """
+        expected_output = """\
+                <?xml version="1.0" encoding="utf-8"?>
+                <root><child1></child1><child2 one="1">2</child2></root>"""
+        new_elements = [
+            NewElement(name="child1", parent_name="root"),
+            NewElement(
+                name="child2", parent_name="root", content="2", attributes={"one": "1"}
+            ),
+        ]
+        self.run_and_assert(new_elements, input_code, expected_output)
+
+    def test_add_new_sibling_same_name(self):
+        input_code = """\
+                <root>
+                    <child>child 1</child>
+                </root>
+        """
+        expected_output = """\
+                <?xml version="1.0" encoding="utf-8"?>
+                <root>
+                    <child>child 1</child>
+                <child>child 2</child></root>"""
+        new_elements = [
+            NewElement(name="child", parent_name="root", content="child 2"),
+        ]
+        self.run_and_assert(new_elements, input_code, expected_output)
+
+    def test_add_nested_elementsl(self):
+        input_code = """\
+        <?xml version="1.0" encoding="utf-8" ?>
+        <configuration>
+          <system.web>
+          </system.web>
+          <system.webServer>
+            <validation validateIntegratedModeConfiguration="false" />
+            <modules>
+              <remove name="ScriptModule" />
+              <add name="ScriptModule" preCondition="managedHandler" type="System.Web.Handlers.ScriptModule, System.Web.Extensions, Version=3.5.0.0, Culture=neutral, PublicKeyToken=31BF3856AD364E35" />
+            </modules>
+          </system.webServer>
+        </configuration>
+        """
+        expected_output = """\
+        <?xml version="1.0" encoding="utf-8"?>
+        <configuration>
+          <system.web>
+          </system.web>
+          <system.webServer>
+            <validation validateIntegratedModeConfiguration="false"></validation>
+            <modules>
+              <remove name="ScriptModule"></remove>
+              <add name="ScriptModule" preCondition="managedHandler" type="System.Web.Handlers.ScriptModule, System.Web.Extensions, Version=3.5.0.0, Culture=neutral, PublicKeyToken=31BF3856AD364E35"></add>
+            </modules>
+          <httpProtocol><customHeaders><add name="X-Frame-Options" value="DENY"></add></customHeaders></httpProtocol></system.webServer>
+        </configuration>"""
+        new_elements = [
+            NewElement(
+                name="httpProtocol",
+                parent_name="system.webServer",
+                content=NewElement(
+                    name="customHeaders",
+                    parent_name="httpProtocol",
+                    content=NewElement(
+                        name="add",
+                        parent_name="customHeaders",
+                        attributes={"name": "X-Frame-Options", "value": "DENY"},
+                    ),
+                ),
+            ),
+        ]
+        self.run_and_assert(new_elements, input_code, expected_output)