Preserve kwarg

pixee · Nov 14, 2023 · 1c5a498 · 1c5a498
1 parent f743d0e
commit 1c5a498
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/src/core_codemods/harden_pyyaml.py b/src/core_codemods/harden_pyyaml.py
@@ -60,6 +60,8 @@ def on_result_found(self, original_node, updated_node):
             self.add_needed_import(self._module_name)
         new_args = [
             *updated_node.args[:1],
-            self.parse_expression(f"{maybe_name}.SafeLoader"),
+            updated_node.args[1].with_changes(
+                value=self.parse_expression(f"{maybe_name}.SafeLoader")
+            ),
         ]
         return self.update_arg_target(updated_node, new_args)
diff --git a/tests/codemods/test_harden_pyyaml.py b/tests/codemods/test_harden_pyyaml.py
@@ -42,7 +42,7 @@ def test_all_unsafe_loaders_kwarg(self, tmpdir, loader):
 
         expected = """import yaml
 data = b'!!python/object/apply:subprocess.Popen \\n- ls'
-deserialized_data = yaml.load(data, yaml.SafeLoader)
+deserialized_data = yaml.load(data, Loader=yaml.SafeLoader)
 """
         self.run_and_assert(tmpdir, input_code, expected)
 
@@ -57,7 +57,7 @@ def test_import_alias(self, tmpdir):
 from yaml import Loader
 
 data = b'!!python/object/apply:subprocess.Popen \\n- ls'
-deserialized_data = yam.load(data, yam.SafeLoader)
+deserialized_data = yam.load(data, Loader=yam.SafeLoader)
 """
         self.run_and_assert(tmpdir, input_code, expected)