From 069a602bdbb1d019656bf43abcd6fd48c833d109 Mon Sep 17 00:00:00 2001
From: Marc Planelles <marc.planelles@ansys.com>
Date: Wed, 11 Sep 2024 15:33:28 +0200
Subject: [PATCH] tweak: small typing improvement (#81)

---
 src/ansys/simai/core/data/types.py     |  4 +++-
 src/ansys/simai/core/utils/requests.py | 18 +++++++++++++++++-
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/ansys/simai/core/data/types.py b/src/ansys/simai/core/data/types.py
index 9beafd98..bc2082df 100644
--- a/src/ansys/simai/core/data/types.py
+++ b/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.
diff --git a/src/ansys/simai/core/utils/requests.py b/src/ansys/simai/core/utils/requests.py
index 1bb4b235..f6c75b30 100644
--- a/src/ansys/simai/core/utils/requests.py
+++ b/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.