return raw response (#340)

Author: jean-lucas

nucleus/__init__.py

@@ -1021,6 +1021,7 @@ def make_request(
         payload: Optional[dict],
         route: str,
         requests_command=requests.post,
+        return_raw_response: bool = False,
     ) -> dict:
         """Makes a request to a Nucleus API endpoint.
 
@@ -1030,9 +1031,10 @@ def make_request(
             payload: Given request payload.
             route: Route for the request.
             Requests command: ``requests.post``, ``requests.get``, or ``requests.delete``.
+            return_raw_response: return the request's response object entirely
 
         Returns:
-            Response payload as JSON dict.
+            Response payload as JSON dict or request object.
         """
         if payload is None:
             payload = {}
@@ -1042,18 +1044,7 @@ def make_request(
                     "Received defined payload with GET request! Will ignore payload"
                 )
             payload = None
-        return self._connection.make_request(payload, route, requests_command)  # type: ignore
-
-    def handle_bad_response(
-        self,
-        endpoint,
-        requests_command,
-        requests_response=None,
-        aiohttp_response=None,
-    ):
-        self._connection.handle_bad_response(
-            endpoint, requests_command, requests_response, aiohttp_response
-        )
+        return self._connection.make_request(payload, route, requests_command, return_raw_response)  # type: ignore
 
     def _set_api_key(self, api_key):
         """Fetch API key from environment variable NUCLEUS_API_KEY if not set"""

nucleus/connection.py

@@ -40,7 +40,11 @@ def put(self, payload: dict, route: str):
         return self.make_request(payload, route, requests_command=requests.put)
 
     def make_request(
-        self, payload: dict, route: str, requests_command=requests.post
+        self,
+        payload: dict,
+        route: str,
+        requests_command=requests.post,
+        return_raw_response: bool = False,
     ) -> dict:
         """
         Makes a request to Nucleus endpoint and logs a warning if not
@@ -49,6 +53,7 @@ def make_request(
         :param payload: given payload
         :param route: route for the request
         :param requests_command: requests.post, requests.get, requests.delete
+        :param return_raw_response: return the request's response object entirely
         :return: response JSON
         """
         endpoint = f"{self.endpoint}/{route}"
@@ -73,6 +78,9 @@ def make_request(
         if not response.ok:
             self.handle_bad_response(endpoint, requests_command, response)
 
+        if return_raw_response:
+            return response
+
         return response.json()
 
     def handle_bad_response(

nucleus/model.py

@@ -234,16 +234,17 @@ def add_tags(self, tags: List[str]):
         Args:
             tags: list of tag names
         """
-        response = self._client.make_request(
+        response: requests.Response = self._client.make_request(
             {MODEL_TAGS_KEY: tags},
             f"model/{self.id}/tag",
             requests_command=requests.post,
+            return_raw_response=True,
         )
 
-        if response.get("msg", False):
+        if response.ok:
             self.tags.extend(tags)
 
-        return response
+        return response.json()
 
     def remove_tags(self, tags: List[str]):
         """Remove tag(s) from the model. ::
@@ -257,13 +258,14 @@ def remove_tags(self, tags: List[str]):
         Args:
             tags: list of tag names to remove
         """
-        response = self._client.make_request(
+        response: requests.Response = self._client.make_request(
             {MODEL_TAGS_KEY: tags},
             f"model/{self.id}/tag",
             requests_command=requests.delete,
+            return_raw_response=True,
         )
 
-        if response.get("msg", False):
+        if response.ok:
             self.tags = list(filter(lambda t: t not in tags, self.tags))
 
-        return response
+        return response.json()