diff --git a/src/core_codemods/harden_pyyaml.py b/src/core_codemods/harden_pyyaml.py index 2bec53956..d35236102 100644 --- 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 index d783e2f19..d5619b46e 100644 --- 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)