update after review

pixee · Sep 27, 2023 · 1e2924b · 1e2924b
1 parent c2f3972
commit 1e2924b
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 24 deletions.
diff --git a/src/codemodder/codemods/api/helpers.py b/src/codemodder/codemods/api/helpers.py
@@ -64,24 +64,30 @@ def replace_arg(
 
         for arg in original_node.args:
             if matchers.matches(arg.keyword, matchers.Name(target_arg_name)):
-                new = cst.Arg(
-                    keyword=cst.parse_expression(target_arg_name),
-                    value=cst.parse_expression(target_arg_replacement_val),
-                    equal=arg.equal,
+                new = self.make_new_arg(
+                    target_arg_name, target_arg_replacement_val, arg
                 )
                 arg_added = True
             else:
                 new = arg
             new_args.append(new)
 
         if add_if_missing and not arg_added:
-            new = cst.Arg(
-                keyword=cst.parse_expression(target_arg_name),
-                value=cst.parse_expression(target_arg_replacement_val),
-                equal=cst.AssignEqual(
-                    whitespace_before=cst.SimpleWhitespace(""),
-                    whitespace_after=cst.SimpleWhitespace(""),
-                ),
-            )
+            new = self.make_new_arg(target_arg_name, target_arg_replacement_val)
             new_args.append(new)
         return new_args
+
+    def make_new_arg(self, name, value, existing_arg=None):
+        equal = (
+            existing_arg.equal
+            if existing_arg
+            else cst.AssignEqual(
+                whitespace_before=cst.SimpleWhitespace(""),
+                whitespace_after=cst.SimpleWhitespace(""),
+            )
+        )
+        return cst.Arg(
+            keyword=cst.parse_expression(name),
+            value=cst.parse_expression(value),
+            equal=equal,
+        )
diff --git a/src/core_codemods/docs/pixee_python_enable-jinja2-autoescape.md b/src/core_codemods/docs/pixee_python_enable-jinja2-autoescape.md
@@ -8,9 +8,9 @@ Our codemod currently checks if you forgot to turn autoescape on or if you expli
   from jinja2 import Environment
 
 - env = Environment()
-- env = Environment(autoescape=False, loader=some-loader)
+- env = Environment(autoescape=False, loader=some_loader)
 + env = Environment(autoescape=True)
-+ env = Environment(autoescape=True, loader=some-loader)
++ env = Environment(autoescape=True, loader=some_loader)
   ...
 ```
 

diff --git a/src/core_codemods/enable_jinja2_autoescape.py b/src/core_codemods/enable_jinja2_autoescape.py
@@ -13,18 +13,11 @@ def rule(cls):
         return """
             rules:
                 - patterns:
-                  - pattern-either:
-                    - pattern: |
-                        jinja2.Environment()
-                    - pattern: |
-                        jinja2.Environment(...)
-                    - pattern: |
-                        jinja2.Environment(..., autoescape=False, ...)
-                  - pattern-not: |
-                      jinja2.Environment(..., autoescape=True, ...)
+                  - pattern: jinja2.Environment(...)
+                  - pattern-not: jinja2.Environment(..., autoescape=True, ...)
                   - pattern-inside: |
                       import jinja2
-                        ...
+                      ...
         """
 
     def on_result_found(self, original_node, updated_node):