Adds dict access to all responses

Author: nunomaduro

algoliasearch/responses.py

@@ -1,6 +1,6 @@
 import abc
 
-from typing import List
+from typing import List, Union
 
 from algoliasearch.exceptions import RequestException
 from algoliasearch.helpers import get_items
@@ -26,6 +26,12 @@ def wait(self):
 
         pass  # pragma: no cover
 
+    @abc.abstractmethod
+    def __getitem__(self, key):
+        # type:(Union[int, str]) -> Union[int, str, dict]
+
+        pass  # pragma: no cover
+
 
 class IndexingResponse(Response):
 
@@ -48,6 +54,11 @@ def wait(self):
 
         return self
 
+    def __getitem__(self, key):
+        # type:(Union[int, str]) -> Union[int, str, dict]
+
+        return self.raw_responses[key]
+
 
 class MultipleResponse(Response):
 
@@ -57,6 +68,9 @@ def __init__(self, responses=None):
         self.responses = [] if responses is None else responses
         self._waitable = list(self.responses)
 
+    def __repr__(self):
+        return self.responses
+
     def push(self, response):
         # type: (Response) -> None
 
@@ -75,6 +89,11 @@ def wait(self):
 
         return self
 
+    def __getitem__(self, key):
+        # type:(Union[int, str]) -> Union[int, str, dict]
+
+        return self.responses[key]
+
 
 class AddApiKeyResponse(Response):
 
@@ -98,6 +117,11 @@ def wait(self):
 
         return self
 
+    def __getitem__(self, key):
+        # type:(Union[int, str]) -> Union[int, str, dict]
+
+        return self.raw_response[key]
+
 
 class UpdateApiKeyResponse(Response):
 
@@ -114,7 +138,8 @@ def wait(self):
 
         while not self._done:
             api_key = self._client._sync().get_api_key(
-                self.raw_response['key']
+                self.raw_response['key'],
+                self._request_options
             )
 
             if self._have_changed(api_key):
@@ -136,6 +161,11 @@ def _have_changed(self, api_key):
         return any([(valid_key in body and body[valid_key] == api_key.get(
             valid_key)) for valid_key in valid_keys])
 
+    def __getitem__(self, key):
+        # type:(Union[int, str]) -> Union[int, str, dict]
+
+        return self.raw_response[key]
+
 
 class DeleteApiKeyResponse(Response):
 
@@ -158,6 +188,11 @@ def wait(self):
 
         return self
 
+    def __getitem__(self, key):
+        # type:(Union[int, str]) -> Union[int, str, dict]
+
+        return self.raw_response[key]
+
 
 class RestoreApiKeyResponse(Response):
 
@@ -182,6 +217,11 @@ def wait(self):
 
         return self
 
+    def __getitem__(self, key):
+        # type:(Union[int, str]) -> Union[int, str, dict]
+
+        return self.raw_response[key]
+
 
 class MultipleIndexBatchIndexingResponse(Response):
 
@@ -201,3 +241,8 @@ def wait(self):
             self._done = True
 
         return self
+
+    def __getitem__(self, key):
+        # type:(Union[int, str]) -> Union[int, str, dict]
+
+        return self.raw_response[key]

tests/unit/test_responses.py

@@ -3,7 +3,11 @@
 import mock
 
 from algoliasearch.exceptions import RequestException
-from algoliasearch.responses import IndexingResponse
+from algoliasearch.responses import (
+    IndexingResponse, MultipleResponse,
+    AddApiKeyResponse, UpdateApiKeyResponse,
+    DeleteApiKeyResponse, MultipleIndexBatchIndexingResponse
+)
 from algoliasearch.search_client import SearchClient
 
 
@@ -18,9 +22,26 @@ def test_dict_access(self):
             'foo': 'bar',
         }
 
-        index = self.client.init_index('foo')
-        response = IndexingResponse(index, [response])
-        self.assertEqual(response.raw_responses[0]['foo'], 'bar')
+        response_object = IndexingResponse({}, [response])
+        self.assertEqual(response_object[0]['foo'], 'bar')
+
+        response_object = MultipleResponse([response])
+        self.assertEqual(response_object[0]['foo'], 'bar')
+
+        response_object = MultipleResponse([IndexingResponse({}, [response])])
+        self.assertEqual(response_object[0][0]['foo'], 'bar')
+
+        response_object = AddApiKeyResponse({}, response)
+        self.assertEqual(response_object['foo'], 'bar')
+
+        response_object = UpdateApiKeyResponse({}, response, {})
+        self.assertEqual(response_object['foo'], 'bar')
+
+        response_object = DeleteApiKeyResponse({}, response, '')
+        self.assertEqual(response_object['foo'], 'bar')
+
+        response_object = MultipleIndexBatchIndexingResponse({}, response)
+        self.assertEqual(response_object['foo'], 'bar')
 
     def test_request_exception(self):
         self.client._transporter.write.side_effect = RequestException('', 300)