Handle enum properties which have None / null as a possible value

For enums which only allow None, add new null_enum.py.jinja template
For enums with mixed values, add logic to remove None from list
Mark enum property as nullable if None isn't the only value
Add tests for enums with mixed values and only null values
TODO: Make the type checker stop failing, but PTO next week

Author: juspence

end_to_end_tests/golden-record/my_test_api_client/api/tests/get_user_list.py

@@ -6,6 +6,8 @@
 from ...client import Client
 from ...models.a_model import AModel
 from ...models.an_enum import AnEnum
+from ...models.an_enum_with_null import AnEnumWithNull
+from ...models.an_enum_with_only_null import AnEnumWithOnlyNull
 from ...models.http_validation_error import HTTPValidationError
 from ...types import UNSET, Response
 
@@ -14,6 +16,8 @@ def _get_kwargs(
     *,
     client: Client,
     an_enum_value: List[AnEnum],
+    an_enum_value_with_null: List[Optional[AnEnumWithNull]],
+    an_enum_value_with_only_null: List[AnEnumWithOnlyNull],
     some_date: Union[datetime.date, datetime.datetime],
 ) -> Dict[str, Any]:
     url = "{}/tests/".format(client.base_url)
@@ -27,13 +31,29 @@ def _get_kwargs(
 
         json_an_enum_value.append(an_enum_value_item)
 
+    json_an_enum_value_with_null = []
+    for an_enum_value_with_null_item_data in an_enum_value_with_null:
+        an_enum_value_with_null_item = (
+            an_enum_value_with_null_item_data.value if an_enum_value_with_null_item_data else None
+        )
+
+        json_an_enum_value_with_null.append(an_enum_value_with_null_item)
+
+    json_an_enum_value_with_only_null = []
+    for an_enum_value_with_only_null_item_data in an_enum_value_with_only_null:
+        an_enum_value_with_only_null_item = an_enum_value_with_only_null_item_data.value
+
+        json_an_enum_value_with_only_null.append(an_enum_value_with_only_null_item)
+
     if isinstance(some_date, datetime.date):
         json_some_date = some_date.isoformat()
     else:
         json_some_date = some_date.isoformat()
 
     params: Dict[str, Any] = {
         "an_enum_value": json_an_enum_value,
+        "an_enum_value_with_null": json_an_enum_value_with_null,
+        "an_enum_value_with_only_null": json_an_enum_value_with_only_null,
         "some_date": json_some_date,
     }
     params = {k: v for k, v in params.items() if v is not UNSET and v is not None}
@@ -82,11 +102,15 @@ def sync_detailed(
     *,
     client: Client,
     an_enum_value: List[AnEnum],
+    an_enum_value_with_null: List[Optional[AnEnumWithNull]],
+    an_enum_value_with_only_null: List[AnEnumWithOnlyNull],
     some_date: Union[datetime.date, datetime.datetime],
 ) -> Response[Union[HTTPValidationError, List[AModel]]]:
     kwargs = _get_kwargs(
         client=client,
         an_enum_value=an_enum_value,
+        an_enum_value_with_null=an_enum_value_with_null,
+        an_enum_value_with_only_null=an_enum_value_with_only_null,
         some_date=some_date,
     )
 
@@ -101,13 +125,17 @@ def sync(
     *,
     client: Client,
     an_enum_value: List[AnEnum],
+    an_enum_value_with_null: List[Optional[AnEnumWithNull]],
+    an_enum_value_with_only_null: List[AnEnumWithOnlyNull],
     some_date: Union[datetime.date, datetime.datetime],
 ) -> Optional[Union[HTTPValidationError, List[AModel]]]:
     """Get a list of things"""
 
     return sync_detailed(
         client=client,
         an_enum_value=an_enum_value,
+        an_enum_value_with_null=an_enum_value_with_null,
+        an_enum_value_with_only_null=an_enum_value_with_only_null,
         some_date=some_date,
     ).parsed
 
@@ -116,11 +144,15 @@ async def asyncio_detailed(
     *,
     client: Client,
     an_enum_value: List[AnEnum],
+    an_enum_value_with_null: List[Optional[AnEnumWithNull]],
+    an_enum_value_with_only_null: List[AnEnumWithOnlyNull],
     some_date: Union[datetime.date, datetime.datetime],
 ) -> Response[Union[HTTPValidationError, List[AModel]]]:
     kwargs = _get_kwargs(
         client=client,
         an_enum_value=an_enum_value,
+        an_enum_value_with_null=an_enum_value_with_null,
+        an_enum_value_with_only_null=an_enum_value_with_only_null,
         some_date=some_date,
     )
 
@@ -134,6 +166,8 @@ async def asyncio(
     *,
     client: Client,
     an_enum_value: List[AnEnum],
+    an_enum_value_with_null: List[Optional[AnEnumWithNull]],
+    an_enum_value_with_only_null: List[AnEnumWithOnlyNull],
     some_date: Union[datetime.date, datetime.datetime],
 ) -> Optional[Union[HTTPValidationError, List[AModel]]]:
     """Get a list of things"""
@@ -142,6 +176,8 @@ async def asyncio(
         await asyncio_detailed(
             client=client,
             an_enum_value=an_enum_value,
+            an_enum_value_with_null=an_enum_value_with_null,
+            an_enum_value_with_only_null=an_enum_value_with_only_null,
             some_date=some_date,
         )
     ).parsed

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

@@ -7,6 +7,8 @@
 from .all_of_sub_model_type_enum import AllOfSubModelTypeEnum
 from .an_all_of_enum import AnAllOfEnum
 from .an_enum import AnEnum
+from .an_enum_with_null import AnEnumWithNull
+from .an_enum_with_only_null import AnEnumWithOnlyNull
 from .an_int_enum import AnIntEnum
 from .another_all_of_sub_model import AnotherAllOfSubModel
 from .another_all_of_sub_model_type import AnotherAllOfSubModelType

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

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

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

@@ -0,0 +1,8 @@
+from enum import Enum
+
+
+class AnEnumWithOnlyNull(Enum):
+    VALUE_0 = None
+
+    def __str__(self) -> str:
+        return str(self.value)

end_to_end_tests/openapi.json

@@ -27,6 +27,30 @@
             "name": "an_enum_value",
             "in": "query"
           },
+          {
+            "required": true,
+            "schema": {
+              "title": "An Enum Value With Null And String Values",
+              "type": "array",
+              "items": {
+                "$ref": "#/components/schemas/AnEnumWithNull"
+              }
+            },
+            "name": "an_enum_value_with_null",
+            "in": "query"
+          },
+          {
+            "required": true,
+            "schema": {
+              "title": "An Enum Value With Only Null Values",
+              "type": "array",
+              "items": {
+                "$ref": "#/components/schemas/AnEnumWithOnlyNull"
+              }
+            },
+            "name": "an_enum_value_with_only_null",
+            "in": "query"
+          },
           {
             "required": true,
             "schema": {
@@ -1164,6 +1188,22 @@
         ],
         "description": "For testing Enums in all the ways they can be used "
       },
+      "AnEnumWithNull": {
+        "title": "AnEnumWithNull",
+        "enum": [
+          "FIRST_VALUE",
+          "SECOND_VALUE",
+          null
+        ],
+        "description": "For testing Enums with mixed string / null values "
+      },
+      "AnEnumWithOnlyNull": {
+        "title": "AnEnumWithOnlyNull",
+        "enum": [
+          null
+        ],
+        "description": "For testing Enums with only null values "
+      },
       "AnAllOfEnum": {
         "title": "AnAllOfEnum",
         "enum": [

openapi_python_client/__init__.py

@@ -227,6 +227,7 @@ def _build_models(self) -> None:
         models_dir.mkdir()
         models_init = models_dir / "__init__.py"
         imports = []
+        NoneType = type(None)
 
         model_template = self.env.get_template("model.py.jinja")
         for model in self.openapi.models:
@@ -236,11 +237,14 @@ def _build_models(self) -> None:
 
         # Generate enums
         str_enum_template = self.env.get_template("str_enum.py.jinja")
+        null_enum_template = self.env.get_template("null_enum.py.jinja")
         int_enum_template = self.env.get_template("int_enum.py.jinja")
         for enum in self.openapi.enums:
             module_path = models_dir / f"{enum.class_info.module_name}.py"
             if enum.value_type is int:
                 module_path.write_text(int_enum_template.render(enum=enum), encoding=self.file_encoding)
+            elif enum.value_type is NoneType:
+                module_path.write_text(null_enum_template.render(enum=enum), encoding=self.file_encoding)
             else:
                 module_path.write_text(str_enum_template.render(enum=enum), encoding=self.file_encoding)
             imports.append(import_string_from_class(enum.class_info))

openapi_python_client/parser/properties/__init__.py

@@ -332,6 +332,14 @@ def build_enum_property(
                 schemas,
             )
 
+    # Remove None from str / int list, if present, and mark property as nullable
+    # If list only has None, with no str or int, make special None enum instead
+    keys_to_remove = [key for key, value in values.items() if value is None]
+    if keys_to_remove and len(keys_to_remove) < len(values.items()):
+        data.nullable = True
+        for key in keys_to_remove:
+            values.pop(key)
+
     for value in values.values():
         value_type = type(value)
         break

openapi_python_client/templates/null_enum.py.jinja

@@ -0,0 +1,9 @@
+from enum import Enum
+
+class {{ enum.class_info.name }}(Enum):
+    {% for key, value in enum.values.items() %}
+    {{ key }} = {{ value }}
+    {% endfor %}
+ 
+    def __str__(self) -> str:
+        return str(self.value)

tests/test_parser/test_properties/test_init.py

@@ -222,6 +222,8 @@ class TestEnumProperty:
             (True, False, "{}"),
             (False, True, "Union[Unset, None, {}]"),
             (True, True, "Optional[{}]"),
+            (True, False, "Optional[{}]"),
+            (True, False, "{}"),
         ),
     )
     def test_get_type_string(self, mocker, enum_property_factory, required, nullable, expected):
@@ -273,6 +275,37 @@ def test_values_from_list_duplicate(self):
         with pytest.raises(ValueError):
             EnumProperty.values_from_list(data)
 
+    def test_values_from_list_with_null(self):
+        from openapi_python_client.parser.properties import EnumProperty
+
+        data = ["abc", "123", "a23", "1bc", 4, -3, "a Thing WIth spaces", "", "null"]
+
+        result = EnumProperty.values_from_list(data)
+
+        # None / null is removed from result, and result is now Optional[{}]
+        assert result == {
+            "ABC": "abc",
+            "VALUE_1": "123",
+            "A23": "a23",
+            "VALUE_3": "1bc",
+            "VALUE_4": 4,
+            "VALUE_NEGATIVE_3": -3,
+            "A_THING_WITH_SPACES": "a Thing WIth spaces",
+            "VALUE_7": "",
+        }
+
+    def test_values_from_list_with_only_null(self):
+        from openapi_python_client.parser.properties import EnumProperty
+
+        data = ["null"]
+
+        result = EnumProperty.values_from_list(data)
+
+        # None / null is not removed from result since it's the only value
+        assert result == {
+            "VALUE_0": None,
+        }
+
 
 class TestPropertyFromData:
     def test_property_from_data_str_enum(self, enum_property_factory):
@@ -304,6 +337,67 @@ def test_property_from_data_str_enum(self, enum_property_factory):
             "ParentAnEnum": prop,
         }
 
+    def test_property_from_data_str_enum_with_null(self, enum_property_factory):
+        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="AnEnumWithNull", enum=["A", "B", "C", "null"], nullable=False, default="B")
+        name = "my_enum"
+        required = True
+
+        schemas = Schemas(classes_by_name={"AnEnum": existing})
+
+        prop, new_schemas = property_from_data(
+            name=name, required=required, data=data, schemas=schemas, parent_name="parent", config=Config()
+        )
+
+        # None / null is removed from enum, and property is now nullable
+        assert prop == enum_property_factory(
+            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",
+        )
+        assert prop.nullable is True
+        assert schemas != new_schemas, "Provided Schemas was mutated"
+        assert new_schemas.classes_by_name == {
+            "AnEnumWithNull": existing,
+            "ParentAnEnum": prop,
+        }
+
+    def test_property_from_data_null_enum(self, enum_property_factory):
+        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="AnEnumWithOnlyNull", enum=["null"], nullable=False, default=None)
+        name = "my_enum"
+        required = True
+
+        schemas = Schemas(classes_by_name={"AnEnum": existing})
+
+        prop, new_schemas = property_from_data(
+            name=name, required=required, data=data, schemas=schemas, parent_name="parent", config=Config()
+        )
+
+        assert prop == enum_property_factory(
+            name=name,
+            required=required,
+            values={"VALUE_0": None},
+            class_info=Class(name="ParentAnEnum", module_name="parent_an_enum"),
+            value_type=type(None),
+            default=None,
+        )
+        assert prop.nullable is False
+        assert schemas != new_schemas, "Provided Schemas was mutated"
+        assert new_schemas.classes_by_name == {
+            "AnEnumWithOnlyNull": existing,
+            "ParentAnEnum": prop,
+        }
+
     def test_property_from_data_int_enum(self, enum_property_factory):
         from openapi_python_client.parser.properties import Class, EnumProperty, Schemas, property_from_data
         from openapi_python_client.schema import Schema