fix: #20: Fix KeyError: 'title' and related issues

python_client_generator/generate_apis.py

@@ -33,21 +33,66 @@ def resolve_property_default(property: Dict[str, Any]) -> Optional[str]:
 
 
 def get_return_type(responses: Dict[str, Any]) -> Optional[str]:
-    successful_responses = [v for k, v in responses.items() if int(k) >= 200 and int(k) < 300]
-    if len(successful_responses) != 1:
-        raise Exception("Incorrect number of successful responses:", len(successful_responses))
+    def check_if_valid_success_response(key: str) -> bool:
+        if key == "default":
+            return True
 
-    schema = successful_responses[0]["content"]["application/json"]["schema"]
-    if not schema:
+        if key == "2XX":
+            return True
+
+        if int(key) >= 200 and int(key) < 300:
+            return True
+
+        return False
+
+    # Only consider successful responses
+    successful_responses_raw = {
+        k: v for k, v in responses.items() if check_if_valid_success_response(k)
+    }
+
+    if len(successful_responses_raw) == 0:
+        return None
+
+    # Pop the default response if there are multiple successful responses (e.g. 200, 201, 204)
+    if len(successful_responses_raw) > 1:
+        successful_responses_raw.pop("default", None)
+
+    # Map the responses to a list
+    successful_responses = [v for _, v in successful_responses_raw.items()]
+
+    # Not all successful responses have a content key, see: https://spec.openapis.org/oas/v3.0.3#responses-object
+    if "content" not in successful_responses[0]:
+        return None
+
+    content = successful_responses[0]["content"]
+
+    if "application/json" not in content:
+        return None
+
+    schema = successful_responses[0]["content"]["application/json"].get("schema")
+    if schema is None:
         return None
 
-    return sanitize_name(schema["title"])
+    if "type" not in schema:
+        return sanitize_name(schema["title"]) if "title" in schema else None
+
+    if schema["type"] == "array":
+        return f"List[{resolve_type(schema['items'])}]"
+    if schema["type"] == "object":
+        return sanitize_name(schema.get("title", "Dict[str, Any]"))
 
 
 def _get_request_body_params(method: Dict[str, Any]) -> List[Dict[str, Any]]:
-    args = []
+    """
+    Supported media types are:
+    - application/json
+    - multipart/form-data
 
+    Media types' ranges are not supported at the moment (e.g. '*/*', 'application/*', etc).
+    """
+    args = []
     content = method["requestBody"]["content"]
+
     if "application/json" in content:
         schema = content["application/json"]["schema"]
         args.append(
@@ -56,9 +101,14 @@ def _get_request_body_params(method: Dict[str, Any]) -> List[Dict[str, Any]]:
                 "schema": schema,
             }
         )
-    else:
+    elif "multipart/form-data" in content:
+        schema = content["multipart/form-data"].get("schema")
+
+        # If schema is not defined with properties, we can't generate arguments
+        if schema is None or "properties" not in schema:
+            return args
+
         # Create argument for each multipart upload property
-        schema = content["multipart/form-data"]["schema"]
         for k, v in schema["properties"].items():
             args.append({"name": k, "schema": v})
 
@@ -76,9 +126,9 @@ def get_function_args(method: Dict[str, Any]) -> List[Dict[str, Any]]:
     keys = ["path", "query", "header"]
     parameters = method.get("parameters", [])
     for k in keys:
-        params += [p for p in parameters if p["in"] == k and p["required"]]
+        params += [p for p in parameters if p["in"] == k and p.get("required", False)]
     for k in keys:
-        params += [p for p in parameters if p["in"] == k and not p["required"]]
+        params += [p for p in parameters if p["in"] == k and not p.get("required", False)]
 
     # Convert params to args format required for templating
     return [

python_client_generator/generate_models.py

@@ -47,8 +47,41 @@ def _get_schema_references(schema: Dict[str, Any]) -> List[str]:
         return []
     elif schema["type"] == "array":
         return _get_schema_references(schema["items"])
-    elif schema["type"] == "object" or "enum" in schema:
-        return [schema["title"]]
+    elif schema["type"] == "object" or (schema["type"] == "string" and "enum" in schema):
+        # As some nested enums may not have a title, we need to check for it.
+        # This is observed to happen inside the properties of a schema that uses an enum with referencing to another enum schema (raw values instead)  # noqa E501
+        # Example:
+        # "properties": { # properties of an object schema (type is object)
+        #     "status": {
+        #         "type": "string",
+        #         "enum": [
+        #             "active",
+        #             "inactive"
+        #         ]
+        #     }
+        # }
+        # In this case, the enum values are not defined in a schema with a title, so we need to check for it.  # noqa E501
+        # For the case where the enum values are defined in a schema with a title, the title will be used.  # noqa E501
+        # Example:
+        # "schemas": {
+        #     ...,  # other schemas
+        #     "Status": {
+        #         "title": "Status",
+        #         "type": "string",
+        #         "enum": [
+        #             "active",
+        #             "inactive"
+        #         ]
+        #     }
+        # }
+        # And then it will be referenced in the properties like this:
+        # "properties": { # properties of an object schema (type is object)
+        #     "status": {
+        #         "$ref": "#/components/schemas/Status"
+        #     }
+        # }
+
+        return [schema.get("title", "")]
     else:
         return []
 

python_client_generator/main.py

@@ -5,7 +5,11 @@
 
 from pathlib import Path
 
-from python_client_generator.utils import assert_openapi_version, dereference_swagger
+from python_client_generator.utils import (
+    add_schema_title_if_missing,
+    assert_openapi_version,
+    dereference_swagger,
+)
 
 from .generate_apis import generate_apis
 from .generate_base_client import generate_base_client
@@ -32,6 +36,7 @@ def main() -> None:
         swagger = json.load(f)
 
     assert_openapi_version(swagger)
+    add_schema_title_if_missing(swagger["components"]["schemas"])
     dereferenced_swagger = dereference_swagger(swagger, swagger)
 
     # Create root directory

python_client_generator/utils.py

@@ -81,7 +81,28 @@ def resolve_type(schema: Dict[str, Any], depth: int = 0, use_literals: bool = Fa
     elif "type" not in schema:
         return "Any"
     elif schema["type"] == "object":
-        if "properties" in schema:
+        # If a schema has properties and a title, we can use the title as the type
+        # name. Otherwise, we just return a generic Dict[str, Any]
+        # This happens when a schema has an object in the properties that doesn't reference another schema.  # noqa E501
+        # Example:
+        # {
+        #    "Schema_Name": {
+        #        "title": "Schema_Name",
+        #        "type": "object",
+        #        "properties": {
+        #            "property_name": {
+        #                "type": "object",
+        #                "properties": {
+        #                    "nested_property": {
+        #                        "type": "string"
+        #                    }
+        #                }
+        #            }
+        #        }
+        #    }
+        # }
+
+        if "properties" in schema and "title" in schema:
             return sanitize_name(schema["title"])
         else:
             return "Dict[str, Any]"
@@ -111,3 +132,30 @@ def resolve_type(schema: Dict[str, Any], depth: int = 0, use_literals: bool = Fa
 def assert_openapi_version(schema: Dict[str, Any]) -> None:
     if not schema.get("openapi") or semver.Version.parse(schema.get("openapi")).major != 3:  # type: ignore # noqa: E501
         raise UnsupportedOpenAPISpec("OpenAPI file provided is not version 3.x")
+
+
+def add_schema_title_if_missing(schemas: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Add 'title' key to schemas if missing to prevent issues with type resolution.
+    Only adds title to object and enum schemas.
+
+    Args:
+        schemas (Dict[str, Any]): Swagger schemas under components.schemas
+    Returns:
+        Dict[str, Any]: Schemas with 'title' key added if missing
+
+    Raises:
+        ValueError: If schema is missing 'type' key
+    """
+
+    for k, v in schemas.items():
+        if "title" not in v and isinstance(v, dict):
+            schema_type = v.get("type")
+
+            if not schema_type:
+                raise ValueError(f"Schema {k} is missing 'type' key")
+
+            if schema_type == "object" or (schema_type == "string" and "enum" in v):
+                v["title"] = k
+
+    return schemas

tests/test_utils.py

@@ -0,0 +1,59 @@
+import pytest
+import typing as t
+from python_client_generator.utils import add_schema_title_if_missing
+
+
+@pytest.mark.parametrize(
+    "raw_schema, expected_schema",
+    [
+        (
+            {"X": {"type": "string", "enum": ["A", "B"]}},
+            {
+                "X": {
+                    "title": "X",
+                    "type": "string",
+                    "enum": ["A", "B"],
+                }
+            },
+        ),  # Should add title to enum schema
+        (
+            {
+                "X": {
+                    "type": "object",
+                    "properties": {
+                        "a": {"type": "string"},
+                        "b": {"type": "string"},
+                    },
+                }
+            },
+            {
+                "X": {
+                    "title": "X",
+                    "type": "object",
+                    "properties": {
+                        "a": {"type": "string"},
+                        "b": {"type": "string"},
+                    },
+                }
+            },
+        ),  # Should add title to object schema
+        (
+            {
+                "X": {
+                    "title": "X",
+                    "type": "string",
+                    "enum": ["A", "B"],
+                }
+            },
+            {
+                "X": {
+                    "title": "X",
+                    "type": "string",
+                    "enum": ["A", "B"],
+                }
+            },
+        ),  # Shouldn't fail if title is already present
+    ],
+)
+def test_add_schema_title_if_missing(raw_schema: t.Dict[str, t.Any], expected_schema) -> None:
+    assert add_schema_title_if_missing(raw_schema) == expected_schema