🐛 Remove trailing comma from ReplaceImportsCodemod

pydantic · Jun 30, 2023 · ce75aee · ce75aee
1 parent 0cd8d34
commit ce75aee
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 48 deletions.
diff --git a/bump_pydantic/codemods/replace_imports.py b/bump_pydantic/codemods/replace_imports.py
@@ -105,6 +105,7 @@ def leave_replace_import(self, _: cst.ImportFrom, updated_node: cst.ImportFrom)
                 AddImportsVisitor.add_needed_import(self.context, *import_info.to_import_str)
                 if len(updated_node.names) > 1:  # type: ignore
                     names = [alias for alias in aliases if alias.name.value != import_info.to_import_str[-1]]
+                    names[-1] = names[-1].with_changes(comma=cst.MaybeSentinel.DEFAULT)
                     updated_node = updated_node.with_changes(names=names)
                 else:
                     return cst.RemoveFromParent()  # type: ignore[return-value]

diff --git a/tests/integration/test_cli.py b/tests/integration/test_cli.py
@@ -129,15 +129,6 @@ def before() -> Folder:
                 "        ...",
             ],
         ),
-        # File(
-        #     "rename_method.py",
-        #     content=[
-        #         "from project.add_none import A",
-        #         "",
-        #         'a = A(a=1, b=2, c=3, d=4, e={"ha": "ha"})',
-        #         "a.dict()",
-        #     ],
-        # ),
         File(
             "replace_generic.py",
             content=[
@@ -206,21 +197,21 @@ def before() -> Folder:
                 "    a: A = Field(A.a, const=True)",
             ],
         ),
-        # File(
-        #     "config_dict_and_settings.py",
-        #     content=[
-        #         "from pydantic import BaseModel, BaseSettings",
-        #         "",
-        #         "",
-        #         "class Settings(BaseSettings):",
-        #         "    sentry_dsn: str",
-        #         "",
-        #         "",
-        #         "class A(BaseModel):",
-        #         "    class Config:",
-        #         "        orm_mode = True",
-        #     ]
-        # )
+        File(
+            "config_dict_and_settings.py",
+            content=[
+                "from pydantic import BaseModel, BaseSettings",
+                "",
+                "",
+                "class Settings(BaseSettings):",
+                "    sentry_dsn: str",
+                "",
+                "",
+                "class A(BaseModel):",
+                "    class Config:",
+                "        orm_mode = True",
+            ],
+        ),
     )
 
 
@@ -277,15 +268,6 @@ def expected() -> Folder:
                 "        ...",
             ],
         ),
-        # File(
-        #     "rename_method.py",
-        #     content=[
-        #         "from project.add_none import A",
-        #         "",
-        #         'a = A(a=1, b=2, c=3, d=4, e={"ha": "ha"})',
-        #         "a.dict()",
-        #     ],
-        # ),
         File(
             "replace_generic.py",
             content=[
@@ -356,21 +338,21 @@ def expected() -> Folder:
                 "    a: Literal[A.a] = A.a",
             ],
         ),
-        # File(
-        #     "config_dict_and_settings.py",
-        #     content=[
-        #         "from pydantic import ConfigDict, BaseModel",
-        #         "from pydantic_settings import BaseSettings",
-        #         "",
-        #         "",
-        #         "class Settings(BaseSettings):",
-        #         "    sentry_dsn: str",
-        #         "",
-        #         "",
-        #         "class A(BaseModel):",
-        #         "    model_config = ConfigDict(orm_mode=True)",
-        #     ]
-        # )
+        File(
+            "config_dict_and_settings.py",
+            content=[
+                "from pydantic import ConfigDict, BaseModel",
+                "from pydantic_settings import BaseSettings",
+                "",
+                "",
+                "class Settings(BaseSettings):",
+                "    sentry_dsn: str",
+                "",
+                "",
+                "class A(BaseModel):",
+                "    model_config = ConfigDict(from_attributes=True)",
+            ],
+        ),
     )