fix class generation for nullable enums

Author: eli-bl

end_to_end_tests/baseline_openapi_3.0.json

@@ -1660,7 +1660,9 @@
           "model",
           "nullable_model",
           "one_of_models",
-          "nullable_one_of_models"
+          "nullable_one_of_models",
+          "nullable_enum_as_ref",
+          "nullable_enum_inline"
         ],
         "type": "object",
         "properties": {
@@ -1830,6 +1832,14 @@
               }
             ],
             "nullable": true
+          },
+          "nullable_enum_as_ref": {
+            "$ref": "#/components/schemas/AnEnumWithNull"
+          },
+          "nullable_enum_inline": {
+            "type": "string",
+            "enum": ["FIRST_VALUE", "SECOND_VALUE", null],
+            "nullable": true
           }
         },
         "description": "A Model for testing all the ways custom objects can be used ",
@@ -1850,6 +1860,7 @@
           "SECOND_VALUE",
           null
         ],
+        "nullable": true,
         "description": "For testing Enums with mixed string / null values "
       },
       "AnEnumWithOnlyNull": {

end_to_end_tests/baseline_openapi_3.1.yaml

@@ -1650,7 +1650,9 @@ info:
         "model",
         "nullable_model",
         "one_of_models",
-        "nullable_one_of_models"
+        "nullable_one_of_models",
+        "nullable_enum_as_ref",
+        "nullable_enum_inline"
       ],
       "type": "object",
       "properties": {
@@ -1840,6 +1842,13 @@ info:
               "$ref": "#/components/schemas/ModelWithUnionProperty"
             }
           ]
+        },
+        "nullable_enum_as_ref": {
+          "$ref": "#/components/schemas/AnEnumWithNull"
+        },
+        "nullable_enum_inline": {
+          "type": ["string", "null"],
+          "enum": ["FIRST_VALUE", "SECOND_VALUE", null]
         }
       },
       "description": "A Model for testing all the ways custom objects can be used ",
@@ -1855,6 +1864,7 @@ info:
     },
     "AnEnumWithNull": {
       "title": "AnEnumWithNull",
+      "type": ["string", "null"],
       "enum": [
         "FIRST_VALUE",
         "SECOND_VALUE",

end_to_end_tests/golden-record/my_test_api_client/models/__init__.py

@@ -4,6 +4,7 @@
 from .a_discriminated_union_type_2 import ADiscriminatedUnionType2
 from .a_form_data import AFormData
 from .a_model import AModel
+from .a_model_nullable_enum_inline import AModelNullableEnumInline
 from .a_model_with_properties_reference_that_are_not_object import AModelWithPropertiesReferenceThatAreNotObject
 from .all_of_has_properties_but_no_type import AllOfHasPropertiesButNoType
 from .all_of_has_properties_but_no_type_type_enum import AllOfHasPropertiesButNoTypeTypeEnum
@@ -91,6 +92,7 @@
     "AllOfSubModel",
     "AllOfSubModelTypeEnum",
     "AModel",
+    "AModelNullableEnumInline",
     "AModelWithPropertiesReferenceThatAreNotObject",
     "AnAllOfEnum",
     "AnArrayWithACircularRefInItemsObjectAdditionalPropertiesAItem",

end_to_end_tests/golden-record/my_test_api_client/models/a_model.py

@@ -4,8 +4,10 @@
 from attrs import define as _attrs_define
 from dateutil.parser import isoparse
 
+from ..models.a_model_nullable_enum_inline import AModelNullableEnumInline
 from ..models.an_all_of_enum import AnAllOfEnum
 from ..models.an_enum import AnEnum
+from ..models.an_enum_with_null import AnEnumWithNull
 from ..models.different_enum import DifferentEnum
 from ..types import UNSET, Unset
 
@@ -33,6 +35,8 @@ class AModel:
         nullable_one_of_models (Union['FreeFormModel', 'ModelWithUnionProperty', None]):
         model (ModelWithUnionProperty):
         nullable_model (Union['ModelWithUnionProperty', None]):
+        nullable_enum_as_ref (Union[AnEnumWithNull, None]): For testing Enums with mixed string / null values
+        nullable_enum_inline (Union[AModelNullableEnumInline, None]):
         any_value (Union[Unset, Any]):
         an_optional_allof_enum (Union[Unset, AnAllOfEnum]):
         nested_list_of_enums (Union[Unset, List[List[DifferentEnum]]]):
@@ -57,6 +61,8 @@ class AModel:
     nullable_one_of_models: Union["FreeFormModel", "ModelWithUnionProperty", None]
     model: "ModelWithUnionProperty"
     nullable_model: Union["ModelWithUnionProperty", None]
+    nullable_enum_as_ref: Union[AnEnumWithNull, None]
+    nullable_enum_inline: Union[AModelNullableEnumInline, None]
     an_allof_enum_with_overridden_default: AnAllOfEnum = AnAllOfEnum.OVERRIDDEN_DEFAULT
     any_value: Union[Unset, Any] = UNSET
     an_optional_allof_enum: Union[Unset, AnAllOfEnum] = UNSET
@@ -122,6 +128,18 @@ def to_dict(self) -> Dict[str, Any]:
         else:
             nullable_model = self.nullable_model
 
+        nullable_enum_as_ref: Union[None, str]
+        if isinstance(self.nullable_enum_as_ref, AnEnumWithNull):
+            nullable_enum_as_ref = self.nullable_enum_as_ref.value
+        else:
+            nullable_enum_as_ref = self.nullable_enum_as_ref
+
+        nullable_enum_inline: Union[None, str]
+        if isinstance(self.nullable_enum_inline, AModelNullableEnumInline):
+            nullable_enum_inline = self.nullable_enum_inline.value
+        else:
+            nullable_enum_inline = self.nullable_enum_inline
+
         any_value = self.any_value
 
         an_optional_allof_enum: Union[Unset, str] = UNSET
@@ -199,6 +217,8 @@ def to_dict(self) -> Dict[str, Any]:
                 "nullable_one_of_models": nullable_one_of_models,
                 "model": model,
                 "nullable_model": nullable_model,
+                "nullable_enum_as_ref": nullable_enum_as_ref,
+                "nullable_enum_inline": nullable_enum_inline,
             }
         )
         if any_value is not UNSET:
@@ -342,6 +362,36 @@ def _parse_nullable_model(data: object) -> Union["ModelWithUnionProperty", None]
 
         nullable_model = _parse_nullable_model(d.pop("nullable_model"))
 
+        def _parse_nullable_enum_as_ref(data: object) -> Union[AnEnumWithNull, None]:
+            if data is None:
+                return data
+            try:
+                if not isinstance(data, str):
+                    raise TypeError()
+                componentsschemas_an_enum_with_null = AnEnumWithNull(data)
+
+                return componentsschemas_an_enum_with_null
+            except:  # noqa: E722
+                pass
+            return cast(Union[AnEnumWithNull, None], data)
+
+        nullable_enum_as_ref = _parse_nullable_enum_as_ref(d.pop("nullable_enum_as_ref"))
+
+        def _parse_nullable_enum_inline(data: object) -> Union[AModelNullableEnumInline, None]:
+            if data is None:
+                return data
+            try:
+                if not isinstance(data, str):
+                    raise TypeError()
+                nullable_enum_inline = AModelNullableEnumInline(data)
+
+                return nullable_enum_inline
+            except:  # noqa: E722
+                pass
+            return cast(Union[AModelNullableEnumInline, None], data)
+
+        nullable_enum_inline = _parse_nullable_enum_inline(d.pop("nullable_enum_inline"))
+
         any_value = d.pop("any_value", UNSET)
 
         _an_optional_allof_enum = d.pop("an_optional_allof_enum", UNSET)
@@ -469,6 +519,8 @@ def _parse_not_required_nullable_model(data: object) -> Union["ModelWithUnionPro
             nullable_one_of_models=nullable_one_of_models,
             model=model,
             nullable_model=nullable_model,
+            nullable_enum_as_ref=nullable_enum_as_ref,
+            nullable_enum_inline=nullable_enum_inline,
             any_value=any_value,
             an_optional_allof_enum=an_optional_allof_enum,
             nested_list_of_enums=nested_list_of_enums,

end_to_end_tests/golden-record/my_test_api_client/models/a_model_nullable_enum_inline.py

@@ -0,0 +1,9 @@
+from enum import Enum
+
+
+class AModelNullableEnumInline(str, Enum):
+    FIRST_VALUE = "FIRST_VALUE"
+    SECOND_VALUE = "SECOND_VALUE"
+
+    def __str__(self) -> str:
+        return str(self.value)

openapi_python_client/parser/properties/enum_property.py

@@ -9,7 +9,6 @@
 
 from ... import Config, utils
 from ... import schema as oai
-from ...schema import DataType
 from ..errors import PropertyError
 from .none import NoneProperty
 from .protocol import PropertyProtocol, Value
@@ -43,7 +42,7 @@ class EnumProperty(PropertyProtocol):
     }
 
     @classmethod
-    def build(  # noqa: PLR0911
+    def build(
         cls,
         *,
         data: oai.Schema,
@@ -75,9 +74,10 @@ def build(  # noqa: PLR0911
         # So instead, if null is a possible value, make the property nullable.
         # Mypy is not smart enough to know that the type is right though
         unchecked_value_list = [value for value in enum if value is not None]  # type: ignore
+        allow_null = len(unchecked_value_list) < len(enum)
 
         # It's legal to have an enum that only contains null as a value, we don't bother constructing an enum for that
-        if len(unchecked_value_list) == 0:
+        if len(unchecked_value_list) == 0 and allow_null:
             return (
                 NoneProperty.build(
                     name=name,
@@ -102,21 +102,6 @@ def build(  # noqa: PLR0911
             Union[List[int], List[str]], unchecked_value_list
         )  # We checked this with all the value_types stuff
 
-        if len(value_list) < len(enum):  # Only one of the values was None, that becomes a union
-            data.oneOf = [
-                oai.Schema(type=DataType.NULL),
-                data.model_copy(update={"enum": value_list, "default": data.default}),
-            ]
-            data.enum = None
-            return UnionProperty.build(
-                data=data,
-                name=name,
-                required=required,
-                schemas=schemas,
-                parent_name=parent_name,
-                config=config,
-            )
-
         class_name = data.title or name
         if parent_name:
             class_name = f"{utils.pascal_case(parent_name)}{utils.pascal_case(class_name)}"
@@ -150,8 +135,38 @@ def build(  # noqa: PLR0911
             return checked_default, schemas
         prop = evolve(prop, default=checked_default)
 
+        # Now, if one of the values was None, wrap the type we just made in a union with None. We're
+        # constructing the UnionProperty directly instead of using UnionProperty.build() because we
+        # do *not* want the usual union behavior of creating ThingType1, ThingType2, etc.
+        returned_prop: EnumProperty | UnionProperty | PropertyError
+        if allow_null:
+            none_prop = NoneProperty.build(
+                name=name,
+                required=required,
+                default=None,
+                python_name=prop.python_name,
+                description=None,
+                example=None,
+            )
+            assert not isinstance(none_prop, PropertyError)
+            union_prop = UnionProperty(
+                name=name,
+                required=required,
+                default=checked_default,
+                python_name=prop.python_name,
+                description=data.description,
+                example=data.example,
+                inner_properties=[
+                    none_prop,
+                    prop,
+                ],
+            )
+            returned_prop = union_prop
+        else:
+            returned_prop = prop
+
         schemas = evolve(schemas, classes_by_name={**schemas.classes_by_name, class_info.name: prop})
-        return prop, schemas
+        return returned_prop, schemas
 
     def convert_value(self, value: Any) -> Value | PropertyError | None:
         if value is None or isinstance(value, Value):

tests/test_parser/test_properties/test_enum_property.py

@@ -1,6 +1,8 @@
 import openapi_python_client.schema as oai
 from openapi_python_client.parser.errors import PropertyError
 from openapi_python_client.parser.properties import EnumProperty, Schemas
+from openapi_python_client.parser.properties.none import NoneProperty
+from openapi_python_client.parser.properties.union import UnionProperty
 
 
 def test_conflict(config):
@@ -66,3 +68,22 @@ def test_unsupported_type(config):
     )
 
     assert isinstance(err, PropertyError)
+
+
+def test_nullable_enum(config):
+    data = oai.Schema(
+        type="string",
+        enum=["a", "b", None],
+        nullable=True,
+    )
+    schemas = Schemas()
+
+    p, _ = EnumProperty.build(
+        data=data, name="prop1", required=True, schemas=schemas, parent_name="parent", config=config
+    )
+
+    assert isinstance(p, UnionProperty)
+    assert len(p.inner_properties) == 2  # noqa: PLR2004
+    assert isinstance(p.inner_properties[0], NoneProperty)
+    assert isinstance(p.inner_properties[1], EnumProperty)
+    assert p.inner_properties[1].class_info.name == "ParentProp1"

tests/test_parser/test_properties/test_init.py

@@ -403,14 +403,23 @@ def test_property_from_data_str_enum(self, enum_property_factory, config):
             "ParentAnEnum": prop,
         }
 
+    @pytest.mark.parametrize(
+        "desc,extra_props",
+        [
+            ("3_1_implicit_type", {}),
+            ("3_1_explicit_type", {"type": ["string", "null"]}),
+            ("3_0_implicit_type", {"nullable": True}),
+            ("3_0_explicit_type", {"type": "string", "nullable": True}),
+        ],
+    )
     def test_property_from_data_str_enum_with_null(
-        self, enum_property_factory, union_property_factory, none_property_factory, config
+        self, desc, extra_props, enum_property_factory, union_property_factory, none_property_factory, config
     ):
         from openapi_python_client.parser.properties import Class, Schemas, property_from_data
         from openapi_python_client.schema import Schema
 
         existing = enum_property_factory()
-        data = Schema(title="AnEnum", enum=["A", "B", "C", None], default="B")
+        data = Schema(title="AnEnum", enum=["A", "B", "C", None], default="B", **extra_props)
         name = "my_enum"
         required = True
 
@@ -423,16 +432,16 @@ def test_property_from_data_str_enum_with_null(
         # None / null is removed from enum, and property is now nullable
         assert isinstance(prop, UnionProperty), "Enums with None should be converted to UnionProperties"
         enum_prop = enum_property_factory(
-            name="my_enum_type_1",
+            name=name,
             required=required,
             values={"A": "A", "B": "B", "C": "C"},
             class_info=Class(name="ParentAnEnum", module_name="parent_an_enum"),
             value_type=str,
             default="ParentAnEnum.B",
         )
-        none_property = none_property_factory(name="my_enum_type_0", required=required)
+        none_property = none_property_factory(name=name, required=required)
         assert prop == union_property_factory(
-            name=name, default="ParentAnEnum.B", inner_properties=[none_property, enum_prop]
+            name=name, default="ParentAnEnum.B", inner_properties=[none_property, enum_prop], required=required
         )
         assert schemas != new_schemas, "Provided Schemas was mutated"
         assert new_schemas.classes_by_name == {