diff --git a/bump_pydantic/codemods/replace_config.py b/bump_pydantic/codemods/replace_config.py index 9eef776..9070547 100644 --- a/bump_pydantic/codemods/replace_config.py +++ b/bump_pydantic/codemods/replace_config.py @@ -15,7 +15,6 @@ CHECK_LINK_COMMENT = "# Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information." REMOVED_KEYS = [ - "allow_mutation", "error_msg_templates", "fields", "getter_dict", @@ -38,6 +37,7 @@ "orm_mode": "from_attributes", "schema_extra": "json_schema_extra", "validate_all": "validate_default", + "allow_mutation": "frozen", } EXTRA_ATTRIBUTE = m.Attribute( @@ -164,16 +164,23 @@ def visit_Assign(self, node: cst.Assign) -> None: self.assign_value = node.value def visit_AssignTarget(self, node: cst.AssignTarget) -> None: - if self.inside_config_class: - keyword = RENAMED_KEYS.get(node.target.value, node.target.value) # type: ignore[attr-defined] + if self.inside_config_class and isinstance(node.target, cst.Name): + keyword = RENAMED_KEYS.get(node.target.value, node.target.value) + value: cst.BaseExpression if m.matches(self.assign_value, EXTRA_ATTRIBUTE): value = cst.SimpleString(value=f'"{self.assign_value.attr.value}"') # type: ignore[attr-defined] RemoveImportsVisitor.remove_unused_import(self.context, "pydantic", "Extra") else: - value = self.assign_value # type: ignore[assignment] + value = self.assign_value + if node.target.value == "allow_mutation": + # The `allow_mutation` keyword is the negative of `frozen`. + if m.matches(value, m.Name(value="False")): + value = cst.Name("True") + elif m.matches(value, m.Name(value="True")): + value = cst.Name("False") self.config_args.append( cst.Arg( - keyword=node.target.with_changes(value=keyword), # type: ignore[arg-type] + keyword=node.target.with_changes(value=keyword), value=value, equal=cst.AssignEqual( whitespace_before=cst.SimpleWhitespace(""), diff --git a/tests/unit/test_replace_config.py b/tests/unit/test_replace_config.py index 8c024a1..e40cb54 100644 --- a/tests/unit/test_replace_config.py +++ b/tests/unit/test_replace_config.py @@ -192,21 +192,37 @@ class Potato(BaseModel): """ self.assertCodemod(before, after) + def test_allow_mutation(self) -> None: + before = """ + from pydantic import BaseModel + + class Potato(BaseModel): + class Config: + allow_mutation = False + """ + after = """ + from pydantic import ConfigDict, BaseModel + + class Potato(BaseModel): + model_config = ConfigDict(frozen=True) + """ + self.assertCodemod(before, after) + def test_removed_keys(self) -> None: before = """ from pydantic import BaseModel class Potato(BaseModel): class Config: - allow_mutation = True + underscore_attrs_are_private = True """ after = """ from pydantic import ConfigDict, BaseModel class Potato(BaseModel): - # TODO[pydantic]: The following keys were removed: `allow_mutation`. + # TODO[pydantic]: The following keys were removed: `underscore_attrs_are_private`. # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information. - model_config = ConfigDict(allow_mutation=True) + model_config = ConfigDict(underscore_attrs_are_private=True) """ self.assertCodemod(before, after) @@ -216,16 +232,16 @@ def test_multiple_removed_keys(self) -> None: class Potato(BaseModel): class Config: - allow_mutation = True + underscore_attrs_are_private = True smart_union = True """ after = """ from pydantic import ConfigDict, BaseModel class Potato(BaseModel): - # TODO[pydantic]: The following keys were removed: `allow_mutation`, `smart_union`. + # TODO[pydantic]: The following keys were removed: `underscore_attrs_are_private`, `smart_union`. # Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information. - model_config = ConfigDict(allow_mutation=True, smart_union=True) + model_config = ConfigDict(underscore_attrs_are_private=True, smart_union=True) """ self.assertCodemod(before, after)