Added Enum support

Author: dbanty

openapi_python_client/__init__.py

@@ -53,3 +53,9 @@ def _build_project(openapi: OpenAPI):
         module_path = models_dir / f"{stringcase.snakecase(schema.title)}.py"
         module_path.write_text(model_template.render(schema=schema))
 
+    # Generate enums
+    enum_template = env.get_template("enum.pyi")
+    for enum in openapi.enums.values():
+        module_path = models_dir / f"{enum.name}.py"
+        module_path.write_text(enum_template.render(enum=enum))
+

openapi_python_client/models/openapi.py

@@ -6,25 +6,28 @@
 
 import stringcase
 
-from .properties import Property, property_from_dict, DateTimeProperty, ListProperty, RefProperty, EnumProperty
+from .properties import Property, property_from_dict, ListProperty, RefProperty, EnumProperty
 
 
 class Method(Enum):
     """ HTTP Methods """
+
     GET = "get"
     POST = "post"
     PATCH = "patch"
 
 
 class ParameterLocation(Enum):
     """ The places Parameters can be put when calling an Endpoint """
+
     QUERY = "query"
     PATH = "path"
 
 
 @dataclass
 class Parameter:
     """ A parameter in an Endpoint """
+
     location: ParameterLocation
     property: Property
 
@@ -33,11 +36,7 @@ def from_dict(d: Dict, /) -> Parameter:
         """ Construct a parameter from it's OpenAPI dict form """
         return Parameter(
             location=ParameterLocation(d["in"]),
-            property=property_from_dict(
-                name=d["name"],
-                required=d["required"],
-                data=d["schema"],
-            ),
+            property=property_from_dict(name=d["name"], required=d["required"], data=d["schema"],),
         )
 
 
@@ -46,6 +45,7 @@ class Endpoint:
     """
     Describes a single endpoint on the server
     """
+
     path: str
     method: Method
     description: Optional[str]
@@ -120,19 +120,25 @@ class OpenAPI:
     security_schemes: Dict
     schemas: Dict[str, Schema]
     endpoints: List[Endpoint]
-    enums: Dict[str, List[str]]
+    enums: Dict[str, EnumProperty]
 
     @staticmethod
     def from_dict(d: Dict, /) -> OpenAPI:
         """ Create an OpenAPI from dict """
         schemas = Schema.dict(d["components"]["schemas"])
-        enums = {}
+        enums: Dict[str, EnumProperty] = {}
         for schema in schemas.values():
             for prop in schema.properties:
-                if isinstance(prop, EnumProperty):
-                    enum_class_name = stringcase.pascalcase(prop.name)
-                    enums[enum_class_name] = prop.values
-                    schema.relative_imports.add(f"from .{prop.name} import {enum_class_name}")
+                if not isinstance(prop, EnumProperty):
+                    continue
+                schema.relative_imports.add(f"from .{prop.name} import {prop.class_name}")
+                if prop.class_name in enums:
+                    # We already have an enum with this name, make sure the values match
+                    assert (
+                        prop.values == enums[prop.class_name].values
+                    ), f"Encountered conflicting enum named {prop.class_name}"
+
+                enums[prop.class_name] = prop
 
         return OpenAPI(
             title=d["info"]["title"],

openapi_python_client/models/properties.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Optional, List, Dict, Union, ClassVar
 import stringcase
 
@@ -40,6 +40,7 @@ def to_string(self) -> str:
 @dataclass
 class DateTimeProperty(Property):
     """ A property of type datetime.datetime """
+
     _type_string: ClassVar[str] = "datetime"
 
     def to_string(self) -> str:
@@ -80,6 +81,7 @@ def to_string(self) -> str:
 @dataclass
 class BooleanProperty(Property):
     """ Property for bool """
+
     _type_string: ClassVar[str] = "bool"
 
     def to_string(self) -> str:
@@ -109,19 +111,38 @@ def to_string(self) -> str:
 class EnumProperty(Property):
     """ A property that should use an enum """
 
-    values: List[str]
+    values: Dict[str, str]
+    class_name: str = field(init=False)
+
+    def __post_init__(self):
+        self.class_name = stringcase.pascalcase(self.name)
 
     def get_type_string(self):
         """ Get a string representation of type that should be used when declaring this property """
-        class_name = stringcase.pascalcase(self.name)
+
         if self.required:
-            return class_name
-        return f"Optional[{class_name}]"
+            return self.class_name
+        return f"Optional[{self.class_name}]"
 
     def to_string(self) -> str:
         """ How this should be declared in a dataclass """
         return f"{self.name}: {self.get_type_string()}"
 
+    @staticmethod
+    def values_from_list(l: List[str], /) -> Dict[str, str]:
+        """ Convert a list of values into dict of {name: value} """
+        output: Dict[str, str] = {}
+
+        for i, value in enumerate(l):
+            if value.isalpha():
+                key = value.upper()
+            else:
+                key = f"VALUE_{i}"
+            assert key not in output, f"Duplicate key {key} in Enum"
+            output[key] = value
+
+        return output
+
 
 @dataclass
 class RefProperty(Property):
@@ -165,23 +186,23 @@ def property_from_dict(
 ) -> Property:
     """ Generate a Property from the OpenAPI dictionary representation of it """
     if "enum" in data:
-        return EnumProperty(name=name, required=required, values=data["enum"],)
+        return EnumProperty(name=name, required=required, values=EnumProperty.values_from_list(data["enum"]))
     if "$ref" in data:
         ref = data["$ref"].split("/")[-1]
-        return RefProperty(name=name, required=required, ref=ref,)
+        return RefProperty(name=name, required=required, ref=ref)
     if data["type"] == "string":
         if "format" not in data:
             return StringProperty(
                 name=name, default=data.get("default"), required=required, pattern=data.get("pattern"),
             )
         elif data["format"] == "date-time":
-            return DateTimeProperty(name=name, required=required,)
+            return DateTimeProperty(name=name, required=required)
     elif data["type"] == "number":
-        return FloatProperty(name=name, default=data.get("default"), required=required,)
+        return FloatProperty(name=name, default=data.get("default"), required=required)
     elif data["type"] == "integer":
-        return IntProperty(name=name, default=data.get("default"), required=required,)
+        return IntProperty(name=name, default=data.get("default"), required=required)
     elif data["type"] == "boolean":
-        return BooleanProperty(name=name, required=required,)
+        return BooleanProperty(name=name, required=required)
     elif data["type"] == "array":
         ref = None
         if "$ref" in data["items"]:
@@ -191,5 +212,5 @@ def property_from_dict(
             _type = _openapi_types_to_python_type_strings[data["items"]["type"]]
         return ListProperty(name=name, required=required, type=_type, ref=ref)
     elif data["type"] == "object":
-        return DictProperty(name=name, required=required,)
+        return DictProperty(name=name, required=required)
     raise ValueError(f"Did not recognize type of {data}")

openapi_python_client/templates/enum.pyi

@@ -0,0 +1,6 @@
+from enum import Enum
+
+class {{ enum.class_name }}(Enum):
+    {% for key, value in enum.values.items() %}
+    {{ key }} = "{{ value }}"
+    {% endfor %}