tweak: small typing improvement (#81)

src/ansys/simai/core/data/types.py

@@ -85,7 +85,9 @@
 """
 
 
-APIResponse = Union[Response, Dict[str, Any], List[Dict[str, Any]]]
+JSON = Union[Dict[str, Any], List[Dict[str, Any]]]
+APIResponse = Union[Response, JSON]
+
 
 MonitorCallback = Callable[[int], None]
 """Callback used to monitor the download or upload of a file.

src/ansys/simai/core/utils/requests.py

@@ -22,10 +22,11 @@
 
 import logging
 from json.decoder import JSONDecodeError
+from typing import Literal, overload
 
 import requests
 
-from ansys.simai.core.data.types import APIResponse
+from ansys.simai.core.data.types import JSON, APIResponse
 from ansys.simai.core.errors import ApiClientError, NotFoundError
 
 logger = logging.getLogger(__name__)
@@ -76,6 +77,21 @@ def handle_http_errors(response: requests.Response) -> None:
             raise ApiClientError(f"{response.status_code}: {response.reason}", response) from e
 
 
+@overload
+def handle_response(response: requests.Response, return_json: Literal[True]) -> JSON:
+    ...
+
+
+@overload
+def handle_response(response: requests.Response, return_json: Literal[False]) -> requests.Response:
+    ...
+
+
+@overload
+def handle_response(response: requests.Response, return_json: bool) -> APIResponse:
+    ...
+
+
 def handle_response(response: requests.Response, return_json: bool = True) -> APIResponse:
     """Handle HTTP errors and return the relevant data from the response.