Added support for enum schemas (#122)

emann · dbanty · commit b16ed73f50bb · 2020-08-05T10:28:10.000-04:00
* Added support for enum schemas. Closes #102 Co-authored-by: Ethan Mann <emann@triaxtec.com>
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -13,10 +13,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Additions
 - The generator can now handle many more errors gracefully, skipping the things it can't generate and continuing
     with the pieces it can.
+- Support for Enums declared in "components/schemas" and references to them. (#120)
     
 ### Internal Changes
 - Switched OpenAPI document parsing to use Pydantic based on a vendored version of 
     [openapi-schema-pydantic](https://github.com/kuimono/openapi-schema-pydantic/) (#103).
+- Tests can now be run on Windows.
 
 
 
diff --git a/end_to_end_tests/fastapi_app/openapi.json b/end_to_end_tests/fastapi_app/openapi.json
@@ -41,10 +41,7 @@
                             "title": "An Enum Value",
                             "type": "array",
                             "items": {
-                                "enum": [
-                                    "FIRST_VALUE",
-                                    "SECOND_VALUE"
-                                ]
+                                "$ref": "#/components/schemas/AnEnum"
                             }
                         },
                         "name": "an_enum_value",
@@ -150,22 +147,15 @@
                 "type": "object",
                 "properties": {
                     "an_enum_value": {
-                        "title": "An Enum Value",
-                        "enum": [
-                            "FIRST_VALUE",
-                            "SECOND_VALUE"
-                        ]
+                        "$ref": "#/components/schemas/AnEnum"
                     },
                     "nested_list_of_enums": {
                         "title": "Nested List Of Enums",
                         "type": "array",
                         "items": {
                             "type": "array",
                             "items": {
-                                "enum": [
-                                    "DIFFERENT",
-                                    "OTHER"
-                                ]
+                                "$ref": "#/components/schemas/DifferentEnum"
                             }
                         },
                         "default": []
@@ -199,6 +189,14 @@
                 },
                 "description": "A Model for testing all the ways custom objects can be used "
             },
+            "AnEnum": {
+                "title": "AnEnum",
+                "enum": [
+                    "FIRST_VALUE",
+                    "SECOND_VALUE"
+                ],
+                "description": "For testing Enums in all the ways they can be used "
+            },
             "Body_upload_file_tests_upload_post": {
                 "title": "Body_upload_file_tests_upload_post",
                 "required": [
@@ -213,6 +211,14 @@
                     }
                 }
             },
+            "DifferentEnum": {
+                "title": "DifferentEnum",
+                "enum": [
+                    "DIFFERENT",
+                    "OTHER"
+                ],
+                "description": "An enumeration."
+            },
             "HTTPValidationError": {
                 "title": "HTTPValidationError",
                 "type": "object",
diff --git a/end_to_end_tests/golden-master/my_test_api_client/api/users.py b/end_to_end_tests/golden-master/my_test_api_client/api/users.py
@@ -7,13 +7,13 @@
 from ..client import AuthenticatedClient, Client
 from ..errors import ApiResponseError
 from ..models.a_model import AModel
-from ..models.an_enum_value import AnEnumValue
+from ..models.an_enum import AnEnum
 from ..models.body_upload_file_tests_upload_post import BodyUploadFileTestsUploadPost
 from ..models.http_validation_error import HTTPValidationError
 
 
 def get_user_list(
-    *, client: Client, an_enum_value: List[AnEnumValue], some_date: Union[date, datetime],
+    *, client: Client, an_enum_value: List[AnEnum], some_date: Union[date, datetime],
 ) -> Union[
     List[AModel], HTTPValidationError,
 ]:
diff --git a/end_to_end_tests/golden-master/my_test_api_client/async_api/users.py b/end_to_end_tests/golden-master/my_test_api_client/async_api/users.py
@@ -7,13 +7,13 @@
 from ..client import AuthenticatedClient, Client
 from ..errors import ApiResponseError
 from ..models.a_model import AModel
-from ..models.an_enum_value import AnEnumValue
+from ..models.an_enum import AnEnum
 from ..models.body_upload_file_tests_upload_post import BodyUploadFileTestsUploadPost
 from ..models.http_validation_error import HTTPValidationError
 
 
 async def get_user_list(
-    *, client: Client, an_enum_value: List[AnEnumValue], some_date: Union[date, datetime],
+    *, client: Client, an_enum_value: List[AnEnum], some_date: Union[date, datetime],
 ) -> Union[
     List[AModel], HTTPValidationError,
 ]:
diff --git a/end_to_end_tests/golden-master/my_test_api_client/models/__init__.py b/end_to_end_tests/golden-master/my_test_api_client/models/__init__.py
@@ -1,9 +1,9 @@
 """ Contains all the data models used in inputs/outputs """
 
 from .a_model import AModel
-from .an_enum_value import AnEnumValue
-from .an_enum_value1 import AnEnumValue1
+from .an_enum import AnEnum
 from .body_upload_file_tests_upload_post import BodyUploadFileTestsUploadPost
+from .different_enum import DifferentEnum
 from .http_validation_error import HTTPValidationError
 from .types import *
 from .validation_error import ValidationError
diff --git a/end_to_end_tests/golden-master/my_test_api_client/models/a_model.py b/end_to_end_tests/golden-master/my_test_api_client/models/a_model.py
@@ -4,19 +4,19 @@
 from datetime import date, datetime
 from typing import Any, Dict, List, Optional, Union, cast
 
-from .an_enum_value import AnEnumValue
-from .an_enum_value1 import AnEnumValue1
+from .an_enum import AnEnum
+from .different_enum import DifferentEnum
 
 
 @dataclass
 class AModel:
     """ A Model for testing all the ways custom objects can be used  """
 
-    an_enum_value: AnEnumValue
+    an_enum_value: AnEnum
     a_camel_date_time: Union[datetime, date]
     a_date: date
-    nested_list_of_enums: Optional[List[List[AnEnumValue1]]] = field(
-        default_factory=lambda: cast(Optional[List[List[AnEnumValue1]]], [])
+    nested_list_of_enums: Optional[List[List[DifferentEnum]]] = field(
+        default_factory=lambda: cast(Optional[List[List[DifferentEnum]]], [])
     )
     some_dict: Optional[Dict[Any, Any]] = field(default_factory=lambda: cast(Optional[Dict[Any, Any]], {}))
 
@@ -56,7 +56,7 @@ def to_dict(self) -> Dict[str, Any]:
 
     @staticmethod
     def from_dict(d: Dict[str, Any]) -> AModel:
-        an_enum_value = AnEnumValue(d["an_enum_value"])
+        an_enum_value = AnEnum(d["an_enum_value"])
 
         def _parse_a_camel_date_time(data: Dict[str, Any]) -> Union[datetime, date]:
             a_camel_date_time: Union[datetime, date]
@@ -78,7 +78,7 @@ def _parse_a_camel_date_time(data: Dict[str, Any]) -> Union[datetime, date]:
         for nested_list_of_enums_item_data in d.get("nested_list_of_enums") or []:
             nested_list_of_enums_item = []
             for nested_list_of_enums_item_item_data in nested_list_of_enums_item_data:
-                nested_list_of_enums_item_item = AnEnumValue1(nested_list_of_enums_item_item_data)
+                nested_list_of_enums_item_item = DifferentEnum(nested_list_of_enums_item_item_data)
 
                 nested_list_of_enums_item.append(nested_list_of_enums_item_item)
 
diff --git a/end_to_end_tests/golden-master/my_test_api_client/models/an_enum.py b/end_to_end_tests/golden-master/my_test_api_client/models/an_enum.py
@@ -1,6 +1,6 @@
 from enum import Enum
 
 
-class AnEnumValue(str, Enum):
+class AnEnum(str, Enum):
     FIRST_VALUE = "FIRST_VALUE"
     SECOND_VALUE = "SECOND_VALUE"
diff --git a/end_to_end_tests/golden-master/my_test_api_client/models/different_enum.py b/end_to_end_tests/golden-master/my_test_api_client/models/different_enum.py
@@ -1,6 +1,6 @@
 from enum import Enum
 
 
-class AnEnumValue1(str, Enum):
+class DifferentEnum(str, Enum):
     DIFFERENT = "DIFFERENT"
     OTHER = "OTHER"
diff --git a/openapi_python_client/parser/openapi.py b/openapi_python_client/parser/openapi.py
@@ -208,16 +208,14 @@ class Model:
     relative_imports: Set[str]
 
     @staticmethod
-    def from_data(*, data: Union[oai.Reference, oai.Schema], name: str) -> Union[Model, ParseError]:
+    def from_data(*, data: oai.Schema, name: str) -> Union[Model, ParseError]:
         """ A single Model from its OAI data
 
         Args:
             data: Data of a single Schema
             name: Name by which the schema is referenced, such as a model name.
                 Used to infer the type name if a `title` property is not available.
         """
-        if isinstance(data, oai.Reference):
-            return ParseError(data=data, detail="Reference schemas are not supported.")
         required_set = set(data.required or [])
         required_properties: List[Property] = []
         optional_properties: List[Property] = []
@@ -258,6 +256,18 @@ def build(*, schemas: Dict[str, Union[oai.Reference, oai.Schema]]) -> Schemas:
         """ Get a list of Schemas from an OpenAPI dict """
         result = Schemas()
         for name, data in schemas.items():
+            if isinstance(data, oai.Reference):
+                result.errors.append(ParseError(data=data, detail="Reference schemas are not supported."))
+                continue
+            if data.enum is not None:
+                EnumProperty(
+                    name=name,
+                    title=data.title or name,
+                    required=True,
+                    default=data.default,
+                    values=EnumProperty.values_from_list(data.enum),
+                )
+                continue
             s = Model.from_data(data=data, name=name)
             if isinstance(s, ParseError):
                 result.errors.append(s)
diff --git a/openapi_python_client/parser/properties.py b/openapi_python_client/parser/properties.py
@@ -263,6 +263,11 @@ def get_all_enums() -> Dict[str, EnumProperty]:
         """ Get all the EnumProperties that have been registered keyed by class name """
         return _existing_enums
 
+    @staticmethod
+    def get_enum(name: str) -> Optional[EnumProperty]:
+        """ Get all the EnumProperties that have been registered keyed by class name """
+        return _existing_enums.get(name)
+
     def get_type_string(self) -> str:
         """ Get a string representation of type that should be used when declaring this property """
 
@@ -304,7 +309,12 @@ class RefProperty(Property):
 
     reference: Reference
 
-    template: ClassVar[str] = "ref_property.pyi"
+    @property
+    def template(self) -> str:  # type: ignore
+        enum = EnumProperty.get_enum(self.reference.class_name)
+        if enum:
+            return "enum_property.pyi"
+        return "ref_property.pyi"
 
     def get_type_string(self) -> str:
         """ Get a string representation of type that should be used when declaring this property """
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -1,4 +1,4 @@
-from pathlib import PosixPath
+from pathlib import Path
 from unittest.mock import MagicMock
 
 import pytest
@@ -35,8 +35,8 @@ def test_config(mocker, _create_new_client):
     result = runner.invoke(app, [f"--config={config_path}", "generate", f"--path={path}"], catch_exceptions=False)
 
     assert result.exit_code == 0
-    load_config.assert_called_once_with(path=PosixPath(config_path))
-    _create_new_client.assert_called_once_with(url=None, path=PosixPath(path))
+    load_config.assert_called_once_with(path=Path(config_path))
+    _create_new_client.assert_called_once_with(url=None, path=Path(path))
 
 
 def test_bad_config(mocker, _create_new_client):
@@ -50,7 +50,7 @@ def test_bad_config(mocker, _create_new_client):
 
     assert result.exit_code == 2
     assert "Unable to parse config" in result.stdout
-    load_config.assert_called_once_with(path=PosixPath(config_path))
+    load_config.assert_called_once_with(path=Path(config_path))
     _create_new_client.assert_not_called()
 
 
@@ -87,7 +87,7 @@ def test_generate_path(self, _create_new_client):
         result = runner.invoke(app, ["generate", f"--path={path}"])
 
         assert result.exit_code == 0
-        _create_new_client.assert_called_once_with(url=None, path=PosixPath(path))
+        _create_new_client.assert_called_once_with(url=None, path=Path(path))
 
     def test_generate_handle_errors(self, _create_new_client):
         _create_new_client.return_value = [GeneratorError(detail="this is a message")]
@@ -166,4 +166,4 @@ def test_update_path(self, _update_existing_client):
         result = runner.invoke(app, ["update", f"--path={path}"])
 
         assert result.exit_code == 0
-        _update_existing_client.assert_called_once_with(url=None, path=PosixPath(path))
+        _update_existing_client.assert_called_once_with(url=None, path=Path(path))
diff --git a/tests/test_openapi_parser/test_openapi.py b/tests/test_openapi_parser/test_openapi.py
@@ -105,19 +105,11 @@ def test_from_data_property_parse_error(self, mocker):
 
         assert result == parse_error
 
-    def test_from_data_parse_error_on_reference(self):
-        from openapi_python_client.parser.openapi import Model
-
-        data = oai.Reference.construct()
-        assert Model.from_data(data=data, name="") == ParseError(
-            data=data, detail="Reference schemas are not supported."
-        )
-
 
 class TestSchemas:
     def test_build(self, mocker):
         from_data = mocker.patch(f"{MODULE_NAME}.Model.from_data")
-        in_data = {1: mocker.MagicMock(), 2: mocker.MagicMock(), 3: mocker.MagicMock()}
+        in_data = {"1": mocker.MagicMock(enum=None), "2": mocker.MagicMock(enum=None), "3": mocker.MagicMock(enum=None)}
         schema_1 = mocker.MagicMock()
         schema_2 = mocker.MagicMock()
         error = ParseError()
@@ -132,6 +124,26 @@ def test_build(self, mocker):
             models={schema_1.reference.class_name: schema_1, schema_2.reference.class_name: schema_2,}, errors=[error]
         )
 
+    def test_build_parse_error_on_reference(self):
+        from openapi_python_client.parser.openapi import Schemas
+
+        ref_schema = oai.Reference.construct()
+        in_data = {1: ref_schema}
+        result = Schemas.build(schemas=in_data)
+        assert result.errors[0] == ParseError(data=ref_schema, detail="Reference schemas are not supported.")
+
+    def test_build_enums(self, mocker):
+        from openapi_python_client.parser.openapi import Schemas
+
+        from_data = mocker.patch(f"{MODULE_NAME}.Model.from_data")
+        enum_property = mocker.patch(f"{MODULE_NAME}.EnumProperty")
+        in_data = {"1": mocker.MagicMock(enum=["val1", "val2", "val3"])}
+
+        Schemas.build(schemas=in_data)
+
+        enum_property.assert_called()
+        from_data.assert_not_called()
+
 
 class TestEndpoint:
     def test_parse_request_form_body(self, mocker):
diff --git a/tests/test_openapi_parser/test_properties.py b/tests/test_openapi_parser/test_properties.py
@@ -320,8 +320,28 @@ def test_get_all_enums(self, mocker):
         assert properties.EnumProperty.get_all_enums() == properties._existing_enums
         properties._existing_enums = {}
 
+    def test_get_enum(self):
+        from openapi_python_client.parser import properties
+
+        properties._existing_enums = {"test": "an enum"}
+        assert properties.EnumProperty.get_enum("test") == "an enum"
+        properties._existing_enums = {}
+
 
 class TestRefProperty:
+    def test_template(self, mocker):
+        from openapi_python_client.parser.properties import RefProperty
+
+        ref_property = RefProperty(
+            name="test", required=True, default=None, reference=mocker.MagicMock(class_name="MyRefClass")
+        )
+
+        assert ref_property.template == "ref_property.pyi"
+
+        mocker.patch(f"{MODULE_NAME}.EnumProperty.get_enum", return_value="an enum")
+
+        assert ref_property.template == "enum_property.pyi"
+
     def test_get_type_string(self, mocker):
         from openapi_python_client.parser.properties import RefProperty
 
