From ba371021c1613c96e7b7e686f2f0bf00a34601e4 Mon Sep 17 00:00:00 2001 From: Camillo Lugaresi Date: Sat, 6 Apr 2024 03:09:58 -0700 Subject: [PATCH 1/2] Convert allow_mutation to frozen --- bump_pydantic/codemods/replace_config.py | 8 ++++++- tests/unit/test_replace_config.py | 28 +++++++++++++++++++----- 2 files changed, 29 insertions(+), 7 deletions(-) diff --git a/bump_pydantic/codemods/replace_config.py b/bump_pydantic/codemods/replace_config.py index 9eef776..e73581f 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( @@ -171,6 +171,12 @@ def visit_AssignTarget(self, node: cst.AssignTarget) -> None: RemoveImportsVisitor.remove_unused_import(self.context, "pydantic", "Extra") else: value = self.assign_value # type: ignore[assignment] + 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] 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) From 648be4614ad8313e1e9f90bbf749060d09b7160a Mon Sep 17 00:00:00 2001 From: Camillo Lugaresi Date: Sat, 6 Apr 2024 18:02:46 -0700 Subject: [PATCH 2/2] fix types --- bump_pydantic/codemods/replace_config.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/bump_pydantic/codemods/replace_config.py b/bump_pydantic/codemods/replace_config.py index e73581f..9070547 100644 --- a/bump_pydantic/codemods/replace_config.py +++ b/bump_pydantic/codemods/replace_config.py @@ -164,13 +164,14 @@ 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")): @@ -179,7 +180,7 @@ def visit_AssignTarget(self, node: cst.AssignTarget) -> None: 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(""),