Skip to content

Commit

Permalink
✨ Add note on removed keys
Browse files Browse the repository at this point in the history
  • Loading branch information
Kludex committed Jun 22, 2023
1 parent 75dadb5 commit 1d6d169
Show file tree
Hide file tree
Showing 2 changed files with 57 additions and 7 deletions.
23 changes: 18 additions & 5 deletions bump_pydantic/codemods/replace_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,11 +7,11 @@
from libcst.codemod.visitors import AddImportsVisitor
from libcst.metadata import PositionProvider

UNABLE_REFACTOR_COMMENT = (
"# TODO[pydantic]: We couldn't refactor this class, please create the `model_config` manually."
)
PREFIX_COMMENT = "# TODO[pydantic]: "
REFACTOR_COMMENT = f"{PREFIX_COMMENT}We couldn't refactor this class, please create the `model_config` manually."
REMOVED_KEYS_COMMENT = f"{PREFIX_COMMENT}The following keys were removed: {{keys}}."
INHERIT_CONFIG_COMMENT = f"{PREFIX_COMMENT}The `Config` class inherits from another class, please create the `model_config` manually." # noqa: E501
CHECK_LINK_COMMENT = "# Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information."
INHERIT_CONFIG_COMMENT = "# TODO[pydantic]: The `Config` class inherits from another class, please create the `model_config` manually." # noqa: E501

REMOVED_KEYS = [
"allow_mutation",
Expand Down Expand Up @@ -135,7 +135,7 @@ def leave_config_class(self, original_node: cst.ClassDef, updated_node: cst.Clas
return updated_node.with_changes(
leading_lines=[
*updated_node.leading_lines,
cst.EmptyLine(comment=cst.Comment(value=(UNABLE_REFACTOR_COMMENT))),
cst.EmptyLine(comment=cst.Comment(value=(REFACTOR_COMMENT))),
cst.EmptyLine(comment=cst.Comment(value=(CHECK_LINK_COMMENT))),
]
)
Expand Down Expand Up @@ -206,6 +206,7 @@ def leave_config_class_childless(self, original_node: cst.ClassDef, updated_node
),
)
],
leading_lines=self._leading_lines_from_removed_keys(self.config_args),
)
if m.matches(statement, m.ClassDef(name=m.Name(value="Config")))
else statement
Expand All @@ -214,6 +215,18 @@ def leave_config_class_childless(self, original_node: cst.ClassDef, updated_node
self.config_args = []
return updated_node.with_changes(body=updated_node.body.with_changes(body=body))

@staticmethod
def _leading_lines_from_removed_keys(args: List[cst.Arg]) -> List[cst.EmptyLine]:
removed_keys = [arg.keyword.value for arg in args if arg.keyword.value in REMOVED_KEYS] # type: ignore
if not removed_keys:
return []

formatted_keys = ", ".join(f"`{key}`" for key in removed_keys)
return [
cst.EmptyLine(comment=cst.Comment(value=REMOVED_KEYS_COMMENT.format(keys=formatted_keys))),
cst.EmptyLine(comment=cst.Comment(value=CHECK_LINK_COMMENT)),
]


if __name__ == "__main__":
import textwrap
Expand Down
41 changes: 39 additions & 2 deletions tests/unit/test_replace_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ class Config:
class Potato2(BaseModel):
class Config:
allow_mutation = True
strict = True
"""
after = """
from pydantic import ConfigDict, BaseModel
Expand All @@ -78,7 +78,7 @@ class Potato(BaseModel):
potato = Potato()
class Potato2(BaseModel):
model_config = ConfigDict(allow_mutation=True)
model_config = ConfigDict(strict=True)
"""
self.assertCodemod(before, after)

Expand Down Expand Up @@ -192,3 +192,40 @@ class Potato(BaseModel):
model_config = ConfigDict(extra="allow")
"""
self.assertCodemod(before, after)

def test_removed_keys(self) -> None:
before = """
from pydantic import BaseModel
class Potato(BaseModel):
class Config:
allow_mutation = True
"""
after = """
from pydantic import ConfigDict, BaseModel
class Potato(BaseModel):
# TODO[pydantic]: The following keys were removed: `allow_mutation`.
# Check https://docs.pydantic.dev/dev-v2/migration/#changes-to-config for more information.
model_config = ConfigDict(allow_mutation=True)
"""
self.assertCodemod(before, after)

def test_multiple_removed_keys(self) -> None:
before = """
from pydantic import BaseModel
class Potato(BaseModel):
class Config:
allow_mutation = True
smart_union = True
"""
after = """
from pydantic import ConfigDict, BaseModel
class Potato(BaseModel):
# TODO[pydantic]: The following keys were removed: `allow_mutation`, `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)
"""
self.assertCodemod(before, after)

0 comments on commit 1d6d169

Please sign in to comment.