diff --git a/bump_pydantic/codemods/replace_config.py b/bump_pydantic/codemods/replace_config.py index 3f36e15..164c3c8 100644 --- a/bump_pydantic/codemods/replace_config.py +++ b/bump_pydantic/codemods/replace_config.py @@ -38,7 +38,6 @@ "schema_extra": "json_schema_extra", "validate_all": "validate_default", } -# TODO: The codemod should not replace `Config` in case of removed keys, right? BASE_MODEL_WITH_CONFIG = m.ClassDef( bases=[ @@ -154,9 +153,10 @@ def visit_Assign(self, node: cst.Assign) -> None: 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] self.config_args.append( cst.Arg( - keyword=node.target, # type: ignore[arg-type] + keyword=node.target.with_changes(value=keyword), # type: ignore[arg-type] value=self.assign_value, equal=cst.AssignEqual( whitespace_before=cst.SimpleWhitespace(""), diff --git a/tests/integration/test_cli.py b/tests/integration/test_cli.py index 96fe182..9910157 100644 --- a/tests/integration/test_cli.py +++ b/tests/integration/test_cli.py @@ -208,7 +208,7 @@ def expected() -> Folder: "", "", "class A(BaseModel):", - " model_config = ConfigDict(orm_mode=True, validate_all=True)", + " model_config = ConfigDict(from_attributes=True, validate_default=True)", "", "", "class BaseConfig:", diff --git a/tests/unit/test_replace_config.py b/tests/unit/test_replace_config.py index 5e48091..0b730b1 100644 --- a/tests/unit/test_replace_config.py +++ b/tests/unit/test_replace_config.py @@ -229,3 +229,19 @@ class Potato(BaseModel): model_config = ConfigDict(allow_mutation=True, smart_union=True) """ self.assertCodemod(before, after) + + def test_renamed_keys(self) -> None: + before = """ + from pydantic import BaseModel + + class Potato(BaseModel): + class Config: + orm_mode = True + """ + after = """ + from pydantic import ConfigDict, BaseModel + + class Potato(BaseModel): + model_config = ConfigDict(from_attributes=True) + """ + self.assertCodemod(before, after)