Merge pull request #8 from allenhaozi/feat/response

format response

Author: dingzhaohan

src/bohrium/__init__.py

@@ -1,7 +1,13 @@
 from ._client import Bohrium, AsyncBohrium
-from ._base_client import AsyncAPIClient, SyncAPIClient
+from ._base_client import AsyncAPIClient, SyncAPIClient, BaseClient
 from ._utils._logs import setup_logging as _setup_logging
 
 _setup_logging()
 
-__all__ = ["Bohrium", "AsyncBohrium", "AsyncAPIClient", "SyncAPIClient"]
+__all__ = [
+    "Bohrium",
+    "AsyncBohrium",
+    "AsyncAPIClient",
+    "SyncAPIClient",
+    "BaseClient",
+]

src/bohrium/_base_client.py

@@ -9,10 +9,10 @@
     Generic,
     Literal,
     Mapping,
+    Optional,
     TypeVar,
     Union,
     cast,
-    Optional
 )
 from urllib.parse import urljoin
 
@@ -114,7 +114,7 @@ class BaseClient(Generic[_HttpxClientT]):
     max_retries: int
     timeout: Union[float, Timeout, None]
     _limits: httpx.Limits
-    _version: str | None
+    _version: Union[str, None]
 
     def __init__(
         self,
@@ -151,7 +151,7 @@ def _build_params(self, custom_params) -> dict[str, str]:
         return params
 
     @property
-    def custom_auth(self) -> httpx.Auth | None:
+    def custom_auth(self) -> Union[httpx.Auth, None]:
         return None
 
     @property
@@ -167,9 +167,7 @@ def default_headers(self) -> dict[str, str]:
 
     @property
     def default_params(self) -> dict[str, str]:
-        return {
-            "accessKey": self.access_key
-        }
+        return {"accessKey": self.access_key}
 
     def platform_headers(self) -> Dict[str, str]:
         return platform_headers(self._version)
@@ -187,7 +185,11 @@ def _request(
         merged_params = self._build_params(kwargs.get("params"))
         try:
             return self._client.request(
-                method.upper(), url, json=json, headers=merged_headers, params=merged_params
+                method.upper(),
+                url,
+                json=json,
+                headers=merged_headers,
+                params=merged_params,
             )
         except httpx.TransportError as e:
             logger.error(f"Transport error: {e}.")

src/bohrium/_client.py

@@ -17,7 +17,7 @@ class Bohrium(SyncAPIClient):
 
     # client options
     access_key: str
-    project_id: str | None
+    project_id: Union[str, None]
 
     def __init__(
         self,

src/bohrium/_response.py

@@ -0,0 +1,54 @@
+import httpx
+from typing import Any, Optional, Dict
+
+
+class APIErrorResponse:
+    def __init__(self, message: str, status_code: int):
+        self.message = message
+        self.status_code = status_code
+
+    def to_dict(self) -> Dict[str, Any]:
+        return {"error": self.message, "status_code": self.status_code}
+
+    def __repr__(self) -> str:
+        return f"<APIErrorResponse(status_code={self.status_code}, message='{self.message}')>"
+
+
+class APIResponse:
+    def __init__(self, response: httpx.Response):
+        self._response = response
+        self.status_code = response.status_code
+        self.headers = response.headers
+        self._json = self._parse_json(response)
+
+    def _parse_json(self, response: httpx.Response) -> Optional[Any]:
+        """Parses the JSON content of the response."""
+        try:
+            return response.json()
+        except ValueError:
+            return APIErrorResponse("Invalid JSON response", response.status_code)
+
+    @property
+    def json(self) -> Optional[Any]:
+        """Returns the JSON content of the response or an APIErrorResponse if parsing fails."""
+        return self._json
+
+    def is_success(self) -> bool:
+        """Returns True if the response status code indicates success."""
+        return 200 <= self.status_code < 300
+
+    def raise_for_status(self):
+        """Raises an HTTPError if the response status code indicates an error."""
+        if not self.is_success():
+            raise httpx.HTTPStatusError(
+                f"HTTP Error {self.status_code} for url {self._response.url}",
+                request=self._response.request,
+                response=self._response,
+            )
+
+    def __repr__(self) -> str:
+        return (
+            f"<APIResponse(status_code={self.status_code}, "
+            f"headers={dict(self.headers)}, "
+            f"json={self._json})>"
+        )

src/bohrium/resources/job/job.py

@@ -8,9 +8,11 @@
 from typing import Optional
 
 from ..._resource import AsyncAPIResource, SyncAPIResource
+from ..._response import APIResponse
 from ...types.job.job import JobAddRequest
 from ..tiefblue.tiefblue import Tiefblue
 
+
 log = logging.getLogger(__name__)
 
 
@@ -19,8 +21,11 @@ class Job(SyncAPIResource):
     def detail(self, job_id):
         log.info(f"detail job {job_id}")
         response = self._client.get(f"/openapi/v1/job/{job_id}")
+
+        log.info(response.json())
         log.debug(response)
-        return response.json().get("data")
+        return APIResponse(response).json
+        #return response.json().get("data")
 
     def submit(
         self,

src/bohrium/resources/tiefblue/tiefblue.py

@@ -4,6 +4,7 @@
 import os
 from tqdm import tqdm
 import time
+from typing import Union
 from ..._base_client import SyncAPIClient
 from httpx import URL
 from pprint import pprint
@@ -28,7 +29,7 @@ class Tiefblue:
 
     def __init__(
             self,
-            base_url: str | URL | None = None,
+            base_url: Union[str, URL] = None,
         ) -> None:
 
         if base_url is None: