Skip to content

Commit

Permalink
♻️ Do not convert constr to Annotated (#97)
Browse files Browse the repository at this point in the history
  • Loading branch information
Kludex committed Jul 17, 2023
1 parent 6e1cb4d commit c8ad8ac
Show file tree
Hide file tree
Showing 3 changed files with 12 additions and 4 deletions.
9 changes: 7 additions & 2 deletions bump_pydantic/codemods/con_func.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,10 @@ def leave_ann_assign_constr_call(self, original_node: cst.AnnAssign, updated_nod
func_name = cast(str, annotation.func.attr.value) # type: ignore
type_name = MAP_FUNC_TO_TYPE[func_name]

# TODO: When FastAPI supports Pydantic 2.0.4+, remove the conditional below.
if func_name == "constr":
return updated_node

needed_import = MAP_TYPE_TO_NEEDED_IMPORT.get(type_name)
if needed_import is not None:
AddImportsVisitor.add_needed_import(context=self.context, **needed_import) # type: ignore[arg-type]
Expand All @@ -82,12 +86,13 @@ def leave_ann_assign_constr_call(self, original_node: cst.AnnAssign, updated_nod
annotation = cst.Annotation(annotation=annotated) # type: ignore[assignment]
return updated_node.with_changes(annotation=annotation)

# TODO: When FastAPI supports Pydantic 2.0.4+, remove the comments below.
@m.leave(CONSTR_CALL)
def leave_constr_call(self, original_node: cst.Call, updated_node: cst.Call) -> cst.Call:
self._remove_import(original_node.func)
AddImportsVisitor.add_needed_import(context=self.context, module="pydantic", obj="StringConstraints")
# AddImportsVisitor.add_needed_import(context=self.context, module="pydantic", obj="StringConstraints")
return updated_node.with_changes(
func=cst.Name("StringConstraints"),
# func=cst.Name("StringConstraints"),
args=[
arg if arg.keyword and arg.keyword.value != "regex" else arg.with_changes(keyword=cst.Name("pattern"))
for arg in updated_node.args
Expand Down
4 changes: 2 additions & 2 deletions tests/integration/cases/con_func.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,14 +23,14 @@
expected=File(
"con_func.py",
content=[
"from pydantic import Field, StringConstraints, BaseModel",
"from pydantic import Field, BaseModel, constr",
"from decimal import Decimal",
"from typing import List, Set",
"from typing_extensions import Annotated",
"",
"",
"class Potato(BaseModel):",
" a: Annotated[str, StringConstraints(pattern='[a-z]+')]",
" a: constr(pattern='[a-z]+')",
" b: Annotated[List[int], Field(min_length=1, max_length=10)]",
" c: Annotated[int, Field(gt=0, lt=10)]",
" d: Annotated[bytes, Field(min_length=1, max_length=10)]",
Expand Down
3 changes: 3 additions & 0 deletions tests/unit/test_con_func.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import pytest
from libcst.codemod import CodemodTest

from bump_pydantic.codemods.con_func import ConFuncCallCommand
Expand All @@ -8,6 +9,7 @@ class TestFieldCommand(CodemodTest):

maxDiff = None

@pytest.mark.xfail(reason="Annotated is not supported yet!")
def test_constr_to_annotated(self) -> None:
before = """
from pydantic import BaseModel, constr
Expand All @@ -24,6 +26,7 @@ class Potato(BaseModel):
"""
self.assertCodemod(before, after)

@pytest.mark.xfail(reason="Annotated is not supported yet!")
def test_pydantic_constr_to_annotated(self) -> None:
before = """
import pydantic
Expand Down

0 comments on commit c8ad8ac

Please sign in to comment.