generalize add arg to add args

src/codemodder/codemods/api/helpers.py

@@ -47,11 +47,18 @@ def update_assign_rhs(self, updated_node: cst.Assign, rhs: str):
     def parse_expression(self, expression: str):
         return cst.parse_expression(expression)
 
-    def replace_args(
-        self,
-        original_node,
-        args_info,
-    ):
+    def replace_args(self, original_node, args_info):
+        """
+        Iterate over the args in original_node and replace each arg
+        with any matching arg in `args_info`.
+
+        :param original_node: libcst node with args attribute.
+        :param list args_info: List of tuples.
+            Ex: [("new_arg", "new_arg_value", True/False)]
+            The True/False value should indicate whether the arg should be added
+            if not present in original_node's args.
+        """
+        assert hasattr(original_node, "args")
         new_args = []
 
         for arg in original_node.args:
@@ -70,36 +77,6 @@ def replace_args(
 
         return new_args
 
-    def replace_arg(
-        self,
-        original_node,
-        target_arg_name,
-        target_arg_replacement_val,
-        add_if_missing=False,
-    ):
-        """Given a node, return its args with one arg's value changed.
-
-        If add_if_missing is True, then if target arg is not present, add it.
-        """
-        assert hasattr(original_node, "args")
-        new_args = []
-        arg_added = False
-
-        for arg in original_node.args:
-            if matchers.matches(arg.keyword, matchers.Name(target_arg_name)):
-                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 = 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

src/core_codemods/enable_jinja2_autoescape.py

@@ -21,7 +21,5 @@ def rule(cls):
         """
 
     def on_result_found(self, original_node, updated_node):
-        new_args = self.replace_arg(
-            original_node, "autoescape", "True", add_if_missing=True
-        )
+        new_args = self.replace_args(original_node, [("autoescape", "True", True)])
         return self.update_arg_target(updated_node, new_args)

src/core_codemods/harden_ruamel.py

@@ -27,5 +27,5 @@ def rule(cls):
         """
 
     def on_result_found(self, original_node, updated_node):
-        new_args = self.replace_arg(original_node, "typ", '"safe"')
+        new_args = self.replace_args(original_node, [("typ", '"safe"', False)])
         return self.update_arg_target(updated_node, new_args)

src/core_codemods/jwt_decode_verify.py

@@ -64,18 +64,12 @@ def replace_options_arg(self, node_args):
             new_args.append(new)
         return new_args
 
-    def replace_arg(
-        self,
-        original_node,
-        target_arg_name,
-        target_arg_replacement_val,
-        add_if_missing=False,
-    ):
-        new_args = super().replace_arg(original_node, "verify", "True")
+    def replace_args(self, original_node, args_info):
+        new_args = super().replace_args(original_node, [("verify", "True", False)])
         return self.replace_options_arg(new_args)
 
     def on_result_found(self, original_node, updated_node):
-        new_args = self.replace_arg(original_node, "verify", "True")
+        new_args = self.replace_args(original_node, [("verify", "True", False)])
         return self.update_arg_target(updated_node, new_args)
 
 

src/core_codemods/requests_verify.py

@@ -22,5 +22,5 @@ def rule(cls):
         """
 
     def on_result_found(self, original_node, updated_node):
-        new_args = self.replace_arg(original_node, "verify", "True")
+        new_args = self.replace_args(original_node, [("verify", "True", False)])
         return self.update_arg_target(updated_node, new_args)