Refactor a bit

pydantic · Jul 10, 2023 · e87e8ee · e87e8ee
1 parent 909324f
commit e87e8ee
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/bump_pydantic/codemods/field.py b/bump_pydantic/codemods/field.py
@@ -99,13 +99,17 @@ def leave_field_call(self, original_node: cst.Call, updated_node: cst.Call) -> c
         new_args: List[cst.Arg] = []
         for arg in updated_node.args:
             if m.matches(arg, m.Arg(keyword=m.Name())):
-                args_dict = dict()
-                keyword = RENAMED_KEYWORDS.get(arg.keyword.value, arg.keyword.value)
-                args_dict["keyword"] = arg.keyword.with_changes(value=keyword) # type: ignore
-                # Check if keyword is `allow_mutation` and if so, invert the value.
-                if arg.keyword.value == "allow_mutation":
-                    args_dict["value"] = arg.value.with_changes(value=str(not(arg.value.value == "True")))
-                new_args.append(arg.with_changes(**args_dict))
+                keyword = RENAMED_KEYWORDS.get(arg.keyword.value, arg.keyword.value)  # type: ignore
+                new_arg = arg.with_changes(keyword=arg.keyword.with_changes(value=keyword))  # type: ignore
+                # The `allow_mutation` keyword argument is a special case. It's the negative of `frozen`.
+                if arg.keyword and arg.keyword.value == "allow_mutation":
+                    value = arg.value
+                    if m.matches(arg.value, m.Name(value="False")):
+                        value = cst.Name("True")
+                    elif m.matches(arg.value, m.Name(value="True")):
+                        value = cst.Name("False")
+                    new_arg = new_arg.with_changes(value=value)
+                new_args.append(new_arg)  # type: ignore
             else:
                 new_args.append(arg)
 

diff --git a/tests/unit/test_field.py b/tests/unit/test_field.py
@@ -118,4 +118,4 @@ class Settings(BaseSettings):
             potato: int = Field(..., frozen=True)
             strawberry: int = Field(..., frozen=False)
         """
-        self.assertCodemod(before, after)
+        self.assertCodemod(before, after)