From 1d6d169c00e4b3b073671930c52664f040386d52 Mon Sep 17 00:00:00 2001
From: Marcelo Trylesinski <marcelotryle@gmail.com>
Date: Thu, 22 Jun 2023 11:42:55 +0200
Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20note=20on=20removed=20keys?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 bump_pydantic/codemods/replace_config.py | 23 ++++++++++---
 tests/unit/test_replace_config.py        | 41 ++++++++++++++++++++++--
 2 files changed, 57 insertions(+), 7 deletions(-)

diff --git a/bump_pydantic/codemods/replace_config.py b/bump_pydantic/codemods/replace_config.py
index 0914f74..3f36e15 100644
--- 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
index 2ac42f3..5e48091 100644
--- a/tests/unit/test_replace_config.py
+++ b/tests/unit/test_replace_config.py
@@ -67,7 +67,7 @@ class Config:
 
         class Potato2(BaseModel):
             class Config:
-                allow_mutation = True
+                strict = True
         """
         after = """
         from pydantic import ConfigDict, BaseModel
@@ -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)
 
@@ -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)