more logical way of validating return types

pyarr/request_handler.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Union
+from typing import Any, Optional, Type, Union
 
 import aiohttp
 import requests
@@ -94,7 +94,7 @@ def _get(
                 "Timeout occurred while connecting to API."
             ) from exception
         response = _process_response(res)
-        return self._return(res, dict if isinstance(response, dict) else list)
+        return self._check_type(res, dict if isinstance(response, dict) else list)
 
     def _post(
         self,
@@ -128,7 +128,7 @@ def _post(
                 "Timeout occurred while connecting to API."
             ) from exception
         response = _process_response(res)
-        return self._return(res, dict if isinstance(response, dict) else list)
+        return self._check_type(res, dict if isinstance(response, dict) else list)
 
     def _put(
         self,
@@ -163,7 +163,7 @@ def _put(
             ) from exception
 
         response = _process_response(res)
-        return self._return(res, dict if isinstance(response, dict) else list)
+        return self._check_type(res, dict if isinstance(response, dict) else list)
 
     def _delete(
         self,
@@ -202,19 +202,18 @@ def _delete(
             assert isinstance(response, Response)
         return response
 
-    def _return(self, res: Response, arg1: type) -> Any:
+    def _check_type(self, value: Any, *types: Type) -> Any:
         """Takes the response and asserts its type
 
         Args:
-            res (Response): Response from request
-            arg1 (type): The type that should be asserted
+            value (Any): Response from request
+            types (Type): The types that should be checked for
 
         Returns:
             Any: Many possible return types
         """
-        response = _process_response(res)
-        assert isinstance(response, arg1)
-        return response
+        assert isinstance(value, types)
+        return value
 
 
 class AsyncRequestHandler(RequestHandler):
@@ -272,7 +271,7 @@ async def _get(
                 auth=self.auth,
             )
         )
-        return self._return(response, dict if isinstance(response, dict) else list)
+        return self._check_type(response, dict, list)
 
     async def _post(
         self,
@@ -301,7 +300,7 @@ async def _post(
                 auth=self.auth,
             )
         )
-        return self._return(response, dict if isinstance(response, dict) else list)
+        return self._check_type(response, dict, list)
 
     async def _put(
         self,
@@ -331,7 +330,7 @@ async def _put(
                 auth=self.auth,
             )
         )
-        return self._return(response, dict if isinstance(response, dict) else list)
+        return self._check_type(response, dict, list)
 
     async def _delete(
         self,
@@ -366,19 +365,6 @@ async def _delete(
             assert isinstance(response, aiohttp.ClientResponse)
         return response
 
-    def _return(self, response: aiohttp.ClientResponse, arg1: type) -> Any:
-        """Takes the response and asserts its type
-
-        Args:
-            res (Response): Response from request
-            arg1 (type): The type that should be asserted
-
-        Returns:
-            Any: Many possible return types
-        """
-        assert isinstance(response, arg1)
-        return response
-
 
 def _process_response(
     res: Response,
@@ -468,16 +454,16 @@ async def _aprocess_response(request: _RequestContextManager):
             if res.status == 500:
                 raise PyarrServerError(
                     f"Internal Server Error: {res.json()['message']}",
-                    res.json(),
+                    await res.json(),
                 )
             if res.status == 502:
                 raise PyarrBadGateway("Bad Gateway. Check your server is accessible.")
 
             content_type = res.headers.get("Content-Type", "")
             if "application/json" in content_type:
-                return res.json()
-            else:
-                assert isinstance(res, aiohttp.ClientResponse)
+                return await res.json()
+
+            assert isinstance(res, aiohttp.ClientResponse)
             return res
 
     except HTTPRequestTimeout as exception: