fix: Fix when the Location header is used.

Author: IAmJSD

autogen.py

@@ -20,6 +20,17 @@
 OPENAPI_URL = "https://api.vantage.sh/v2/oas_v3.json"
 OUTPUT_DIR = Path(__file__).parent / "src" / "vantage"
 
+# Maps a substring found in a response description to the internal client method
+# that should handle it, and the Python return type to emit.
+# Checked against each HTTP status code's description during endpoint parsing.
+RESPONSE_HANDLERS: list[tuple[str, str, str]] = [
+    (
+        "will be available at the location specified in the Location header",
+        "_request_for_location",
+        "str",
+    ),
+]
+
 
 @dataclass
 class Parameter:
@@ -48,6 +59,8 @@ class Endpoint:
     request_body_type: str | None = None
     response_type: str | None = None
     is_multipart: bool = False
+    response_handler: str | None = None  # internal client method to call, if not the default
+    response_handler_return_type: str | None = None
 
 
 @dataclass
@@ -292,13 +305,27 @@ def parse_endpoints(schema: dict[str, Any]) -> list[Endpoint]:
 
             response_type = extract_response_type(spec.get("responses", {}), schemas)
 
+            description = spec.get("description")
+
+            response_handler = None
+            response_handler_return_type = None
+            for resp_desc in spec.get("responses", {}).values():
+                text = resp_desc.get("description", "")
+                for phrase, handler, return_type in RESPONSE_HANDLERS:
+                    if phrase in text:
+                        response_handler = handler
+                        response_handler_return_type = return_type
+                        break
+                if response_handler:
+                    break
+
             endpoints.append(
                 Endpoint(
                     path=path,
                     method=method.upper(),
                     operation_id=operation_id,
                     summary=spec.get("summary"),
-                    description=spec.get("description"),
+                    description=description,
                     deprecated=spec.get("deprecated", False),
                     parameters=parameters,
                     request_body_required=request_body.get("required", False)
@@ -307,6 +334,8 @@ def parse_endpoints(schema: dict[str, Any]) -> list[Endpoint]:
                     request_body_type=body_type,
                     response_type=response_type,
                     is_multipart=is_multipart,
+                    response_handler=response_handler,
+                    response_handler_return_type=response_handler_return_type,
                 )
             )
 
@@ -523,6 +552,18 @@ def generate_pydantic_models(schema: dict[str, Any]) -> str:
     return "\n".join(lines)
 
 
+def _collect_handler_routes(resources: dict[str, Resource]) -> dict[str, list[tuple[str, str]]]:
+    """Scan all endpoints and group (method, path) pairs by their response_handler."""
+    handler_routes: dict[str, list[tuple[str, str]]] = {}
+    for resource in resources.values():
+        for endpoint in resource.endpoints:
+            if endpoint.response_handler:
+                handler_routes.setdefault(endpoint.response_handler, []).append(
+                    (endpoint.method, endpoint.path)
+                )
+    return handler_routes
+
+
 def generate_sync_client(resources: dict[str, Resource]) -> str:
     """Generate synchronous client code."""
     lines = [
@@ -618,13 +659,30 @@ def generate_sync_client(resources: dict[str, Resource]) -> str:
             "                body=response.text,",
             "            )",
             "",
+        ]
+    )
+
+    # Inject generated routing: one if-block per handler, checking (method, path)
+    handler_routes = _collect_handler_routes(resources)
+    for handler, routes in sorted(handler_routes.items()):
+        route_set = "{" + ", ".join(f'("{m}", "{p}")' for m, p in sorted(routes)) + "}"
+        lines.append(f"        if (method, path) in {route_set}:")
+        lines.append(f"            return self.{handler}(response)")
+        lines.append("")
+
+    lines.extend(
+        [
             "        try:",
             "            data = response.json()",
             "        except Exception:",
             "            data = None",
             "",
             "        return data",
             "",
+            "    def _request_for_location(self, response: Any) -> str:",
+            '        """Extract the Location header from a response."""',
+            '        return response.headers["Location"]',
+            "",
             "",
         ]
     )
@@ -699,7 +757,10 @@ def generate_sync_method(endpoint: Endpoint, method_name: str) -> list[str]:
 
     # Method signature
     param_str = ", ".join(["self"] + params) if params else "self"
-    return_type = endpoint.response_type or "None"
+    if endpoint.response_handler:
+        return_type = endpoint.response_handler_return_type or "Any"
+    else:
+        return_type = endpoint.response_type or "None"
     lines.append(f"    def {method_name}({param_str}) -> {return_type}:")
 
     # Docstring
@@ -740,7 +801,11 @@ def generate_sync_method(endpoint: Endpoint, method_name: str) -> list[str]:
         lines.append("        body_data = None")
 
     # Make request and coerce response payload into typed models where possible
-    if endpoint.response_type is None:
+    if endpoint.response_handler:
+        lines.append(
+            f'        return self._client.request("{endpoint.method}", path, params=params, body=body_data)'
+        )
+    elif endpoint.response_type is None:
         lines.append(
             f'        self._client.request("{endpoint.method}", path, params=params, body=body_data)'
         )
@@ -849,13 +914,30 @@ def generate_async_client(resources: dict[str, Resource]) -> str:
             "                body=response.text,",
             "            )",
             "",
+        ]
+    )
+
+    # Inject generated routing: one if-block per handler, checking (method, path)
+    handler_routes = _collect_handler_routes(resources)
+    for handler, routes in sorted(handler_routes.items()):
+        route_set = "{" + ", ".join(f'("{m}", "{p}")' for m, p in sorted(routes)) + "}"
+        lines.append(f"        if (method, path) in {route_set}:")
+        lines.append(f"            return self.{handler}(response)")
+        lines.append("")
+
+    lines.extend(
+        [
             "        try:",
             "            data = response.json()",
             "        except Exception:",
             "            data = None",
             "",
             "        return data",
             "",
+            "    def _request_for_location(self, response: Any) -> str:",
+            '        """Extract the Location header from a response."""',
+            '        return response.headers["Location"]',
+            "",
             "",
         ]
     )
@@ -930,7 +1012,10 @@ def generate_async_method(endpoint: Endpoint, method_name: str) -> list[str]:
 
     # Method signature
     param_str = ", ".join(["self"] + params) if params else "self"
-    return_type = endpoint.response_type or "None"
+    if endpoint.response_handler:
+        return_type = endpoint.response_handler_return_type or "Any"
+    else:
+        return_type = endpoint.response_type or "None"
     lines.append(f"    async def {method_name}({param_str}) -> {return_type}:")
 
     # Docstring
@@ -971,7 +1056,11 @@ def generate_async_method(endpoint: Endpoint, method_name: str) -> list[str]:
         lines.append("        body_data = None")
 
     # Make request and coerce response payload into typed models where possible
-    if endpoint.response_type is None:
+    if endpoint.response_handler:
+        lines.append(
+            f'        return await self._client.request("{endpoint.method}", path, params=params, body=body_data)'
+        )
+    elif endpoint.response_type is None:
         lines.append(
             f'        await self._client.request("{endpoint.method}", path, params=params, body=body_data)'
         )

src/vantage/_async/client.py

@@ -123,13 +123,20 @@ async def request(
                 body=response.text,
             )
 
+        if (method, path) in {("POST", "/costs/data_exports"), ("POST", "/kubernetes_efficiency_reports/data_exports"), ("POST", "/unit_costs/data_exports")}:
+            return self._request_for_location(response)
+
         try:
             data = response.json()
         except Exception:
             data = None
 
         return data
 
+    def _request_for_location(self, response: Any) -> str:
+        """Extract the Location header from a response."""
+        return response.headers["Location"]
+
 
 class AccessGrantsAsyncApi:
     """Async API methods for access_grants resource."""
@@ -1111,7 +1118,7 @@ class CostsAsyncApi:
     def __init__(self, client: AsyncClient) -> None:
         self._client = client
 
-    async def create_export(self, body: CreateCostExport, *, groupings: Optional[List[str]] = None) -> None:
+    async def create_export(self, body: CreateCostExport, *, groupings: Optional[List[str]] = None) -> str:
         """
         Generate cost data export
         
@@ -1122,7 +1129,7 @@ async def create_export(self, body: CreateCostExport, *, groupings: Optional[Lis
             "groupings": groupings,
         }
         body_data = body.model_dump(by_alias=True, exclude_none=True) if hasattr(body, 'model_dump') else body
-        await self._client.request("POST", path, params=params, body=body_data)
+        return await self._client.request("POST", path, params=params, body=body_data)
 
     async def list(self, *, cost_report_token: Optional[str] = None, filter: Optional[str] = None, workspace_token: Optional[str] = None, start_date: Optional[str] = None, end_date: Optional[str] = None, groupings: Optional[List[str]] = None, order: Optional[str] = None, limit: Optional[int] = None, page: Optional[int] = None, date_bin: Optional[str] = None, settings_include_credits: Optional[bool] = None, settings_include_refunds: Optional[bool] = None, settings_include_discounts: Optional[bool] = None, settings_include_tax: Optional[bool] = None, settings_amortize: Optional[bool] = None, settings_unallocated: Optional[bool] = None, settings_aggregate_by: Optional[str] = None, settings_show_previous_period: Optional[bool] = None) -> Costs:
         """
@@ -1772,7 +1779,7 @@ async def create(self, body: CreateKubernetesEfficiencyReport) -> KubernetesEffi
             return KubernetesEfficiencyReport.model_validate(data)
         return data
 
-    async def create_export(self, body: CreateKubernetesEfficiencyReportExport, *, groupings: Optional[List[str]] = None) -> DataExport:
+    async def create_export(self, body: CreateKubernetesEfficiencyReportExport, *, groupings: Optional[List[str]] = None) -> str:
         """
         Generate Kubernetes efficiency data export
         
@@ -1783,10 +1790,7 @@ async def create_export(self, body: CreateKubernetesEfficiencyReportExport, *, g
             "groupings": groupings,
         }
         body_data = body.model_dump(by_alias=True, exclude_none=True) if hasattr(body, 'model_dump') else body
-        data = await self._client.request("POST", path, params=params, body=body_data)
-        if isinstance(data, dict):
-            return DataExport.model_validate(data)
-        return data
+        return await self._client.request("POST", path, params=params, body=body_data)
 
     async def get(self, kubernetes_efficiency_report_token: str) -> KubernetesEfficiencyReport:
         """
@@ -2858,7 +2862,7 @@ class UnitCostsAsyncApi:
     def __init__(self, client: AsyncClient) -> None:
         self._client = client
 
-    async def create_export(self, body: CreateUnitCostsExport) -> DataExport:
+    async def create_export(self, body: CreateUnitCostsExport) -> str:
         """
         Generate data export of unit costs
         
@@ -2867,10 +2871,7 @@ async def create_export(self, body: CreateUnitCostsExport) -> DataExport:
         path = "/unit_costs/data_exports"
         params = None
         body_data = body.model_dump(by_alias=True, exclude_none=True) if hasattr(body, 'model_dump') else body
-        data = await self._client.request("POST", path, params=params, body=body_data)
-        if isinstance(data, dict):
-            return DataExport.model_validate(data)
-        return data
+        return await self._client.request("POST", path, params=params, body=body_data)
 
     async def list(self, *, cost_report_token: str, start_date: Optional[str] = None, end_date: Optional[str] = None, date_bin: Optional[str] = None, order: Optional[str] = None, limit: Optional[int] = None, page: Optional[int] = None) -> UnitCosts:
         """

src/vantage/_sync/client.py

@@ -123,13 +123,20 @@ def request(
                 body=response.text,
             )
 
+        if (method, path) in {("POST", "/costs/data_exports"), ("POST", "/kubernetes_efficiency_reports/data_exports"), ("POST", "/unit_costs/data_exports")}:
+            return self._request_for_location(response)
+
         try:
             data = response.json()
         except Exception:
             data = None
 
         return data
 
+    def _request_for_location(self, response: Any) -> str:
+        """Extract the Location header from a response."""
+        return response.headers["Location"]
+
 
 class AccessGrantsApi:
     """API methods for access_grants resource."""
@@ -1111,7 +1118,7 @@ class CostsApi:
     def __init__(self, client: SyncClient) -> None:
         self._client = client
 
-    def create_export(self, body: CreateCostExport, *, groupings: Optional[List[str]] = None) -> None:
+    def create_export(self, body: CreateCostExport, *, groupings: Optional[List[str]] = None) -> str:
         """
         Generate cost data export
         
@@ -1122,7 +1129,7 @@ def create_export(self, body: CreateCostExport, *, groupings: Optional[List[str]
             "groupings": groupings,
         }
         body_data = body.model_dump(by_alias=True, exclude_none=True) if hasattr(body, 'model_dump') else body
-        self._client.request("POST", path, params=params, body=body_data)
+        return self._client.request("POST", path, params=params, body=body_data)
 
     def list(self, *, cost_report_token: Optional[str] = None, filter: Optional[str] = None, workspace_token: Optional[str] = None, start_date: Optional[str] = None, end_date: Optional[str] = None, groupings: Optional[List[str]] = None, order: Optional[str] = None, limit: Optional[int] = None, page: Optional[int] = None, date_bin: Optional[str] = None, settings_include_credits: Optional[bool] = None, settings_include_refunds: Optional[bool] = None, settings_include_discounts: Optional[bool] = None, settings_include_tax: Optional[bool] = None, settings_amortize: Optional[bool] = None, settings_unallocated: Optional[bool] = None, settings_aggregate_by: Optional[str] = None, settings_show_previous_period: Optional[bool] = None) -> Costs:
         """
@@ -1772,7 +1779,7 @@ def create(self, body: CreateKubernetesEfficiencyReport) -> KubernetesEfficiency
             return KubernetesEfficiencyReport.model_validate(data)
         return data
 
-    def create_export(self, body: CreateKubernetesEfficiencyReportExport, *, groupings: Optional[List[str]] = None) -> DataExport:
+    def create_export(self, body: CreateKubernetesEfficiencyReportExport, *, groupings: Optional[List[str]] = None) -> str:
         """
         Generate Kubernetes efficiency data export
         
@@ -1783,10 +1790,7 @@ def create_export(self, body: CreateKubernetesEfficiencyReportExport, *, groupin
             "groupings": groupings,
         }
         body_data = body.model_dump(by_alias=True, exclude_none=True) if hasattr(body, 'model_dump') else body
-        data = self._client.request("POST", path, params=params, body=body_data)
-        if isinstance(data, dict):
-            return DataExport.model_validate(data)
-        return data
+        return self._client.request("POST", path, params=params, body=body_data)
 
     def get(self, kubernetes_efficiency_report_token: str) -> KubernetesEfficiencyReport:
         """
@@ -2858,7 +2862,7 @@ class UnitCostsApi:
     def __init__(self, client: SyncClient) -> None:
         self._client = client
 
-    def create_export(self, body: CreateUnitCostsExport) -> DataExport:
+    def create_export(self, body: CreateUnitCostsExport) -> str:
         """
         Generate data export of unit costs
         
@@ -2867,10 +2871,7 @@ def create_export(self, body: CreateUnitCostsExport) -> DataExport:
         path = "/unit_costs/data_exports"
         params = None
         body_data = body.model_dump(by_alias=True, exclude_none=True) if hasattr(body, 'model_dump') else body
-        data = self._client.request("POST", path, params=params, body=body_data)
-        if isinstance(data, dict):
-            return DataExport.model_validate(data)
-        return data
+        return self._client.request("POST", path, params=params, body=body_data)
 
     def list(self, *, cost_report_token: str, start_date: Optional[str] = None, end_date: Optional[str] = None, date_bin: Optional[str] = None, order: Optional[str] = None, limit: Optional[int] = None, page: Optional[int] = None) -> UnitCosts:
         """