Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Processing extra fields #147

Open
wants to merge 4 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 30 additions & 0 deletions bump_pydantic/codemods/field.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pydantic can't be a dependency of bump_pydantic.

The test passes because the test suit does have pydantic...


RENAMED_KEYWORDS = {
"min_items": "min_length",
Expand Down Expand Up @@ -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:
Expand All @@ -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)


Expand Down
28 changes: 15 additions & 13 deletions bump_pydantic/codemods/replace_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
30 changes: 30 additions & 0 deletions tests/unit/test_field.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)