Processing extra fields

bump_pydantic/codemods/field.py

@@ -4,6 +4,7 @@
 from libcst import matchers as m
 from libcst.codemod import CodemodContext, VisitorBasedCodemodCommand
 from libcst.codemod.visitors import AddImportsVisitor, RemoveImportsVisitor
+from pydantic import Field
 
 RENAMED_KEYWORDS = {
     "min_items": "min_length",
@@ -111,9 +112,27 @@ def leave_field_call(self, original_node: cst.Call, updated_node: cst.Call) -> c
         if not self.has_field_import or not self.inside_field_assign:
             return updated_node
 
+        json_schema_extra_elements: List[cst.DictElement] = []
         new_args: List[cst.Arg] = []
         for arg in updated_node.args:
             if m.matches(arg, m.Arg(keyword=m.Name())):
+                if arg.keyword is not None:
+                    if arg.keyword.value == "json_schema_extra":
+                        json_schema_extra_elements.extend(arg.value.elements)  # type: ignore
+                        continue
+
+                    if (
+                        (arg.keyword.value not in RENAMED_KEYWORDS)
+                        and (arg.keyword.value not in Field.__annotations__)
+                        and (arg.keyword != "extra")
+                    ):
+                        new_dict_element = cst.DictElement(
+                            key=cst.SimpleString(value=f'"{arg.keyword.value}"'),
+                            value=arg.value,
+                        )
+                        json_schema_extra_elements.append(new_dict_element)
+                        continue
+
                 keyword = RENAMED_KEYWORDS.get(arg.keyword.value, arg.keyword.value)  # type: ignore
                 value = arg.value
                 if arg.keyword:
@@ -131,6 +150,17 @@ def leave_field_call(self, original_node: cst.Call, updated_node: cst.Call) -> c
             else:
                 new_args.append(arg)
 
+        if len(json_schema_extra_elements) > 0:
+            extra_arg = cst.Arg(
+                value=cst.Dict(elements=json_schema_extra_elements),
+                keyword=cst.Name(value="json_schema_extra"),
+                equal=cst.AssignEqual(
+                    whitespace_before=cst.SimpleWhitespace(""), whitespace_after=cst.SimpleWhitespace("")
+                ),
+            )
+
+            new_args.append(extra_arg)
+
         return updated_node.with_changes(args=new_args)
 
 

bump_pydantic/codemods/replace_config.py

@@ -217,20 +217,22 @@ def leave_config_class_childless(self, original_node: cst.ClassDef, updated_node
         AddImportsVisitor.add_needed_import(context=self.context, **needed_import)  # type: ignore[arg-type]
         block = cst.ensure_type(updated_node.body, cst.IndentedBlock)
         body = [
-            cst.SimpleStatementLine(
-                body=[
-                    cst.Assign(
-                        targets=[cst.AssignTarget(target=cst.Name("model_config"))],
-                        value=cst.Call(
-                            func=cst.Name("SettingsConfigDict" if self.is_base_settings else "ConfigDict"),
-                            args=self.config_args,
-                        ),
-                    )
-                ],
-                leading_lines=self._leading_lines_from_removed_keys(self.config_args),
+            (
+                cst.SimpleStatementLine(
+                    body=[
+                        cst.Assign(
+                            targets=[cst.AssignTarget(target=cst.Name("model_config"))],
+                            value=cst.Call(
+                                func=cst.Name("SettingsConfigDict" if self.is_base_settings else "ConfigDict"),
+                                args=self.config_args,
+                            ),
+                        )
+                    ],
+                    leading_lines=self._leading_lines_from_removed_keys(self.config_args),
+                )
+                if m.matches(statement, m.ClassDef(name=m.Name(value="Config")))
+                else statement
             )
-            if m.matches(statement, m.ClassDef(name=m.Name(value="Config")))
-            else statement
             for statement in block.body
         ]
         self.is_base_settings = False

tests/unit/test_field.py

@@ -151,3 +151,33 @@ class Settings(BaseSettings):
             potato: int = Field(..., examples=[1])
         """
         self.assertCodemod(before, after)
+
+    def test_json_schema_extra_exist(self) -> None:
+        before = """
+        from pydantic import BaseModel, Field
+
+        class Human(BaseModel):
+            name: str = Field(..., some_extra_field="some_extra_field_value", json_schema_extra={"a": "b"})
+        """
+        after = """
+        from pydantic import BaseModel, Field
+
+        class Human(BaseModel):
+            name: str = Field(..., json_schema_extra={"some_extra_field": "some_extra_field_value", "a": "b"})
+        """
+        self.assertCodemod(before, after)
+
+    def test_json_schema_extra_not_exist(self) -> None:
+        before = """
+        from pydantic import BaseModel, Field
+
+        class Human(BaseModel):
+            name: str = Field(..., min_length=1, some_extra_field="some_extra_field_value")
+        """
+        after = """
+        from pydantic import BaseModel, Field
+
+        class Human(BaseModel):
+            name: str = Field(..., min_length=1, json_schema_extra={"some_extra_field": "some_extra_field_value"})
+        """
+        self.assertCodemod(before, after)