diff --git a/datamodel_code_generator/model/enum.py b/datamodel_code_generator/model/enum.py index 9dce6dec1..c56a05e53 100644 --- a/datamodel_code_generator/model/enum.py +++ b/datamodel_code_generator/model/enum.py @@ -46,6 +46,7 @@ def __init__( description: Optional[str] = None, type_: Optional[Types] = None, default: Any = UNDEFINED, + fallback: Any = UNDEFINED, ): super().__init__( @@ -70,6 +71,9 @@ def __init__( *self.base_classes, ] + if fallback != UNDEFINED: + self.extra_template_data.update({"fallback": fallback}) + @classmethod def get_data_type(cls, types: Types, **kwargs: Any) -> DataType: raise NotImplementedError diff --git a/datamodel_code_generator/model/template/Enum.jinja2 b/datamodel_code_generator/model/template/Enum.jinja2 index dea045499..6ce5e5227 100644 --- a/datamodel_code_generator/model/template/Enum.jinja2 +++ b/datamodel_code_generator/model/template/Enum.jinja2 @@ -15,3 +15,8 @@ class {{ class_name }}({{ base_class }}): """ {%- endif %} {%- endfor -%} +{%- if fallback %} + @classmethod + def _missing_(cls, value): + return {{ class_name }}.{{ fallback }} +{%- endif %} diff --git a/datamodel_code_generator/parser/jsonschema.py b/datamodel_code_generator/parser/jsonschema.py index 72ae9c113..85406c6f4 100644 --- a/datamodel_code_generator/parser/jsonschema.py +++ b/datamodel_code_generator/parser/jsonschema.py @@ -189,6 +189,7 @@ def validate_ref(cls, value: Any) -> Any: ref: Optional[str] = Field(default=None, alias='$ref') nullable: Optional[bool] = False x_enum_varnames: List[str] = Field(default=[], alias='x-enum-varnames') + x_enum_fallback_value: Optional[str] = Field(default=None, alias='x-enum-fallback-value') description: Optional[str] title: Optional[str] example: Any @@ -1015,6 +1016,27 @@ def parse_enum( ) ) + if obj.x_enum_fallback_value is not None: + field_name = str(obj.x_enum_fallback_value) + default = f"'{field_name.translate(escape_characters)}'" + if field_name not in exclude_field_names: + field_name = self.model_resolver.get_valid_field_name( + field_name, model_type=ModelType.ENUM + ) + enum_fields.append( + self.data_model_field_type( + name=field_name, + default=default, + data_type=self.data_type_manager.get_data_type( + Types.any, + ), + required=True, + strip_default_none=self.strip_default_none, + has_default=obj.has_default, + use_field_description=self.use_field_description, + ) + ) + def create_enum(reference_: Reference) -> DataType: enum = Enum( reference=reference_, @@ -1026,6 +1048,7 @@ def create_enum(reference_: Reference) -> DataType: if self.use_subclass_enum and isinstance(obj.type, str) else None, default=obj.default if obj.has_default else UNDEFINED, + fallback=obj.x_enum_fallback_value if obj.x_enum_fallback_value is not None else UNDEFINED, ) self.results.append(enum) return self.data_type(reference=reference_) diff --git a/tests/data/expected/parser/openapi/openapi_parser_x_enum_fallback_value/output.py b/tests/data/expected/parser/openapi/openapi_parser_x_enum_fallback_value/output.py new file mode 100644 index 000000000..88b8448c3 --- /dev/null +++ b/tests/data/expected/parser/openapi/openapi_parser_x_enum_fallback_value/output.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from enum import Enum +from typing import Optional + +from pydantic import BaseModel + + +class EnumTypeWithFallback(Enum): + a = 'a' + b = 'b' + unknown = 'unknown' + + @classmethod + def _missing_(cls, value): + return EnumTypeWithFallback.unknown + + +class EnumTypeWithFallbackExistingField(Enum): + a = 'a' + b = 'b' + unknown = 'unknown' + + @classmethod + def _missing_(cls, value): + return EnumTypeWithFallbackExistingField.unknown + + +class EnumTypeWithFallbackAndDifferentDefault(Enum): + a = 'a' + b = 'b' + unknown = 'unknown' + + @classmethod + def _missing_(cls, value): + return EnumTypeWithFallbackAndDifferentDefault.unknown + + +class EnumTypeWithFallbackAndSameDefault(Enum): + a = 'a' + b = 'b' + unknown = 'unknown' + + @classmethod + def _missing_(cls, value): + return EnumTypeWithFallbackAndSameDefault.unknown + + +class TopLevelModel(BaseModel): + enum_field: EnumTypeWithFallback + enum_field_with_default: Optional[EnumTypeWithFallbackAndDifferentDefault] = 'a' + enum_field_with_default_fallback: Optional[ + EnumTypeWithFallbackAndSameDefault + ] = 'unknown' diff --git a/tests/data/openapi/x_enum_fallback_value.yaml b/tests/data/openapi/x_enum_fallback_value.yaml new file mode 100644 index 000000000..903f713b5 --- /dev/null +++ b/tests/data/openapi/x_enum_fallback_value.yaml @@ -0,0 +1,50 @@ +openapi: 3.0 +components: + schemas: + enum_type_with_fallback: + type: string + description: Enum field with fallback value 'unknown'. + enum: + - 'a' + - 'b' + x-enum-fallback-value: 'unknown' + + enum_type_with_fallback_existing_field: + type: string + description: Enum field with fallback value 'unknown'. + enum: + - 'a' + - 'b' + - 'unknown' + x-enum-fallback-value: 'unknown' + + enum_type_with_fallback_and_different_default: + type: string + description: Enum field with default value 'a' and fallback value 'unknown'. + enum: + - 'a' + - 'b' + - 'unknown' + default: 'a' + x-enum-fallback-value: 'unknown' + + enum_type_with_fallback_and_same_default: + type: string + description: Enum field with default value and fallback value 'unknown'. + enum: + - 'a' + - 'b' + - 'unknown' + default: 'unknown' + x-enum-fallback-value: 'unknown' + + TopLevelModel: + required: + - enum_field + properties: + enum_field: + $ref: "#/components/schemas/enum_type_with_fallback" + enum_field_with_default: + $ref: "#/components/schemas/enum_type_with_fallback_and_different_default" + enum_field_with_default_fallback: + $ref: "#/components/schemas/enum_type_with_fallback_and_same_default" diff --git a/tests/parser/test_openapi.py b/tests/parser/test_openapi.py index 9af6ae3bb..34b979061 100644 --- a/tests/parser/test_openapi.py +++ b/tests/parser/test_openapi.py @@ -379,6 +379,14 @@ class UnknownTypeNumber(Enum): ) +def test_openapi_parser_parse_x_enum_fallback_value(): + parser = OpenAPIParser( + Path(DATA_PATH / 'x_enum_fallback_value.yaml'), + ) + expected_dir = EXPECTED_OPEN_API_PATH / 'openapi_parser_x_enum_fallback_value' + assert parser.parse() == (expected_dir / 'output.py').read_text() + + @pytest.mark.skipif( pydantic.VERSION < '1.9.0', reason='Require Pydantic version 1.9.0 or later ' )