diff --git a/src/datamodel_code_generator/config.py b/src/datamodel_code_generator/config.py index c91d51c52..1c7ab67f9 100644 --- a/src/datamodel_code_generator/config.py +++ b/src/datamodel_code_generator/config.py @@ -44,6 +44,7 @@ DataModel, DataModelFieldBase, ) +from datamodel_code_generator.model.pydantic_v2 import UnionMode # noqa: TC001 - used by Pydantic at runtime from datamodel_code_generator.model.scalar import DataTypeScalar from datamodel_code_generator.model.union import DataTypeUnion from datamodel_code_generator.parser import DefaultPutDict, LiteralType @@ -51,10 +52,6 @@ from datamodel_code_generator.util import ConfigDict, is_pydantic_v2 from datamodel_code_generator.validators import ModelValidators # noqa: TC001 - used by Pydantic at runtime -if TYPE_CHECKING: - from datamodel_code_generator.model.pydantic_v2 import UnionMode - - CallableSchema = Callable[[str], str] DumpResolveReferenceAction = Callable[[Iterable[str]], str] DefaultPutDictSchema = DefaultPutDict[str, str] diff --git a/tests/data/expected/main/generate_config_union_mode.py b/tests/data/expected/main/generate_config_union_mode.py new file mode 100644 index 000000000..99988e6f0 --- /dev/null +++ b/tests/data/expected/main/generate_config_union_mode.py @@ -0,0 +1,10 @@ +# generated by datamodel-codegen: +# filename: + +from __future__ import annotations + +from pydantic import BaseModel, Field + + +class Model(BaseModel): + value: str | int | None = Field(None, union_mode='left_to_right') \ No newline at end of file diff --git a/tests/main/test_main_general.py b/tests/main/test_main_general.py index 4f1e37387..982c666a4 100644 --- a/tests/main/test_main_general.py +++ b/tests/main/test_main_general.py @@ -26,6 +26,7 @@ from datamodel_code_generator.arguments import _dataclass_arguments from datamodel_code_generator.config import GenerateConfig from datamodel_code_generator.format import CodeFormatter, PythonVersion +from datamodel_code_generator.model.pydantic_v2 import UnionMode from datamodel_code_generator.parser.openapi import OpenAPIParser from tests.conftest import assert_output, create_assert_file_content, freeze_time from tests.main.conftest import ( @@ -2062,6 +2063,21 @@ def test_generate_with_config_object(output_file: Path) -> None: assert "user_name" in content +@pytest.mark.skipif(pydantic.VERSION < "2.0.0", reason="GenerateConfig requires Pydantic v2") +def test_generate_config_with_union_mode() -> None: + """Test GenerateConfig with union_mode field.""" + config = GenerateConfig( + output_model_type=DataModelType.PydanticV2BaseModel, + union_mode=UnionMode.left_to_right, + disable_timestamp=True, + ) + result = generate( + input_='{"type": "object", "properties": {"value": {"anyOf": [{"type": "string"}, {"type": "integer"}]}}}', + config=config, + ) + assert_output(result, EXPECTED_MAIN_PATH / "generate_config_union_mode.py") + + @pytest.mark.skipif(pydantic.VERSION < "2.0.0", reason="GenerateConfig requires Pydantic v2") def test_generate_with_config_and_kwargs_raises_error(output_file: Path) -> None: """Test generate() raises error when both config and kwargs are provided."""