✨ Add note on removed keys

pydantic · Jun 22, 2023 · 7a80eff · 7a80eff
1 parent 75dadb5
commit 7a80eff
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 5 deletions.
diff --git a/bump_pydantic/codemods/replace_config.py b/bump_pydantic/codemods/replace_config.py
@@ -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",
@@ -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))),
                 ]
             )
@@ -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
@@ -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

diff --git a/tests/unit/test_replace_config.py b/tests/unit/test_replace_config.py
@@ -192,3 +192,21 @@ 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)