fix import lxml to import lxml.etree

pixee · Oct 6, 2023 · 6cc4453 · 6cc4453
1 parent 4cd522c
commit 6cc4453
Show file tree

Hide file tree

Showing 8 changed files with 19 additions and 19 deletions.
diff --git a/integration_tests/test_lxml_safe_parser_defaults.py b/integration_tests/test_lxml_safe_parser_defaults.py
@@ -11,6 +11,6 @@ class TestLxmlSafeParserDefaults(BaseIntegrationTest):
     original_code, expected_new_code = original_and_expected_from_code_path(
         code_path, [(1, "parser = lxml.etree.XMLParser(resolve_entities=False)\n")]
     )
-    expected_diff = "--- \n+++ \n@@ -1,2 +1,2 @@\n import lxml\n-parser = lxml.etree.XMLParser()\n+parser = lxml.etree.XMLParser(resolve_entities=False)\n"
+    expected_diff = "--- \n+++ \n@@ -1,2 +1,2 @@\n import lxml.etree\n-parser = lxml.etree.XMLParser()\n+parser = lxml.etree.XMLParser(resolve_entities=False)\n"
     expected_line_change = "2"
     change_description = LxmlSafeParserDefaults.CHANGE_DESCRIPTION
diff --git a/integration_tests/test_lxml_safe_parsing.py b/integration_tests/test_lxml_safe_parsing.py
@@ -21,7 +21,7 @@ class TestLxmlSafeParsing(BaseIntegrationTest):
             ),
         ],
     )
-    expected_diff = '--- \n+++ \n@@ -1,3 +1,3 @@\n import lxml\n-lxml.etree.parse("path_to_file")\n-lxml.etree.fromstring("xml_str")\n+lxml.etree.parse("path_to_file", parser=lxml.etree.XMLParser(resolve_entities=False))\n+lxml.etree.fromstring("xml_str", parser=lxml.etree.XMLParser(resolve_entities=False))\n'
+    expected_diff = '--- \n+++ \n@@ -1,3 +1,3 @@\n import lxml.etree\n-lxml.etree.parse("path_to_file")\n-lxml.etree.fromstring("xml_str")\n+lxml.etree.parse("path_to_file", parser=lxml.etree.XMLParser(resolve_entities=False))\n+lxml.etree.fromstring("xml_str", parser=lxml.etree.XMLParser(resolve_entities=False))\n'
     expected_line_change = "2"
     num_changes = 2
     change_description = LxmlSafeParsing.CHANGE_DESCRIPTION
diff --git a/src/core_codemods/docs/pixee_python_safe-lxml-parser-defaults.md b/src/core_codemods/docs/pixee_python_safe-lxml-parser-defaults.md
@@ -11,7 +11,7 @@ Parameter `resolve_entities` has an unsafe default value of `True`. This codemod
 The changes look as follows:
 
 ```diff
-  import lxml
+  import lxml.etree
 
 - parser = lxml.etree.XMLParser()
 - parser = lxml.etree.XMLParser(resolve_entities=True)

diff --git a/src/core_codemods/docs/pixee_python_safe-lxml-parsing.md b/src/core_codemods/docs/pixee_python_safe-lxml-parsing.md
@@ -6,7 +6,7 @@ attacks and external entity (XXE) attacks.
 The changes look as follows:
 
 ```diff
- import lxml
+  import lxml.etree
 - lxml.etree.parse("path_to_file")
 - lxml.etree.fromstring("xml_str")
 + lxml.etree.parse("path_to_file", parser=lxml.etree.XMLParser(resolve_entities=False))

diff --git a/src/core_codemods/lxml_safe_parser_defaults.py b/src/core_codemods/lxml_safe_parser_defaults.py
@@ -27,7 +27,7 @@ def rule(cls):
                           - pattern: XMLTreeBuilder
                           - pattern: XMLPullParser
                   - pattern-inside: |
-                      import lxml
+                      import lxml.etree
                       ...
         """
 

diff --git a/src/core_codemods/lxml_safe_parsing.py b/src/core_codemods/lxml_safe_parsing.py
@@ -26,7 +26,7 @@ def rule(cls):
                             - pattern: parse
                             - pattern: fromstring
                     - pattern-inside: |
-                        import lxml
+                        import lxml.etree
                         ...
                   - patterns:
                     - pattern: lxml.etree.$FUNC(..., parser=None, ...)
@@ -37,13 +37,13 @@ def rule(cls):
                             - pattern: parse
                             - pattern: fromstring
                     - pattern-inside: |
-                        import lxml
+                        import lxml.etree
                         ...
         """
 
     def on_result_found(self, original_node, updated_node):
         self.remove_unused_import(original_node)
-        self.add_needed_import("lxml")
+        self.add_needed_import("lxml.etree")
         safe_parser = "lxml.etree.XMLParser(resolve_entities=False)"
         new_args = self.replace_args(
             original_node,

diff --git a/tests/codemods/test_lxml_safe_parameter_defaults.py b/tests/codemods/test_lxml_safe_parameter_defaults.py
@@ -15,12 +15,12 @@ def test_name(self):
 
     @each_class
     def test_import(self, tmpdir, klass):
-        input_code = f"""import lxml
+        input_code = f"""import lxml.etree
 
 parser = lxml.etree.{klass}()
 var = "hello"
 """
-        expexted_output = f"""import lxml
+        expexted_output = f"""import lxml.etree
 
 parser = lxml.etree.{klass}(resolve_entities=False)
 var = "hello"
@@ -104,11 +104,11 @@ def test_import_alias(self, tmpdir, klass):
     )
     @each_class
     def test_verify_variations(self, tmpdir, klass, input_args, expected_args):
-        input_code = f"""import lxml
+        input_code = f"""import lxml.etree
 parser = lxml.etree.{klass}({input_args})
 var = "hello"
 """
-        expexted_output = f"""import lxml
+        expexted_output = f"""import lxml.etree
 parser = lxml.etree.{klass}({expected_args})
 var = "hello"
 """

diff --git a/tests/codemods/test_lxml_safe_parsing.py b/tests/codemods/test_lxml_safe_parsing.py
@@ -13,12 +13,12 @@ def test_name(self):
 
     @each_func
     def test_import(self, tmpdir, func):
-        input_code = f"""import lxml
+        input_code = f"""import lxml.etree
 
 lxml.etree.{func}("path_to_file")
 var = "hello"
 """
-        expexted_output = f"""import lxml
+        expexted_output = f"""import lxml.etree
 
 lxml.etree.{func}("path_to_file", parser=lxml.etree.XMLParser(resolve_entities=False))
 var = "hello"
@@ -34,7 +34,7 @@ def test_from_import(self, tmpdir, func):
 var = "hello"
 """
         expexted_output = f"""from lxml.etree import {func}
-import lxml
+import lxml.etree
 
 {func}("path_to_file", parser=lxml.etree.XMLParser(resolve_entities=False))
 var = "hello"
@@ -50,7 +50,7 @@ def test_from_import_module(self, tmpdir, func):
 var = "hello"
 """
         expexted_output = f"""from lxml import etree
-import lxml
+import lxml.etree
 
 etree.{func}("path_to_file", parser=lxml.etree.XMLParser(resolve_entities=False))
 var = "hello"
@@ -66,7 +66,7 @@ def test_import_alias(self, tmpdir, func):
 var = "hello"
 """
         expexted_output = f"""from lxml.etree import {func} as func
-import lxml
+import lxml.etree
 
 func("path_to_file", parser=lxml.etree.XMLParser(resolve_entities=False))
 var = "hello"
@@ -98,11 +98,11 @@ def test_import_alias(self, tmpdir, func):
     )
     @each_func
     def test_verify_variations(self, tmpdir, func, input_args, expected_args):
-        input_code = f"""import lxml
+        input_code = f"""import lxml.etree
 lxml.etree.{func}({input_args})
 var = "hello"
 """
-        expexted_output = f"""import lxml
+        expexted_output = f"""import lxml.etree
 lxml.etree.{func}({expected_args})
 var = "hello"
 """