fixes & impl for async

Author: UmanShahzad

ipinfo/handler.py

@@ -116,7 +116,7 @@ def getBatchDetails(
         input list).
 
         The input list is broken up into batches to abide by API requirements.
-        The batch size can be adjusted with `batch_size` but is clipped to 
+        The batch size can be adjusted with `batch_size` but is clipped to
         `BATCH_MAX_SIZE`.
         Defaults to `BATCH_MAX_SIZE`.
 
@@ -158,14 +158,22 @@ def getBatchDetails(
             else:
                 lookup_addresses.append(ip_address)
 
-        # prepare req http options
-        req_opts = {**self.request_options, "timeout": timeout_per_batch}
+        # all in cache - return early.
+        if len(lookup_addresses) == 0:
+            return result
 
+        # do start timer if necessary
         if timeout_total is not None:
             start_time = time.time()
 
+        # prepare req http options
+        req_opts = {**self.request_options, "timeout": timeout_per_batch}
+
         # loop over batch chunks and do lookup for each.
-        for i in range(0, len(ip_addresses), batch_size):
+        url = API_URL + "/batch"
+        headers = handler_utils.get_headers(self.access_token)
+        headers["content-type"] = "application/json"
+        for i in range(0, len(lookup_addresses), batch_size):
             # quit if total timeout is reached.
             if (
                 timeout_total is not None
@@ -176,14 +184,11 @@ def getBatchDetails(
                 else:
                     return result
 
-            chunk = ip_addresses[i : i + batch_size]
+            chunk = lookup_addresses[i : i + batch_size]
 
             # lookup
-            url = API_URL + "/batch"
-            headers = handler_utils.get_headers(self.access_token)
-            headers["content-type"] = "application/json"
             response = requests.post(
-                url, json=lookup_addresses, headers=headers, **req_opts
+                url, json=chunk, headers=headers, **req_opts
             )
 
             # fail on bad status codes

ipinfo/handler_async.py

@@ -118,7 +118,14 @@ async def getDetails(self, ip_address=None):
 
         return Details(details)
 
-    async def getBatchDetails(self, ip_addresses, batch_size=None):
+    async def getBatchDetails(
+        self,
+        ip_addresses,
+        batch_size=None,
+        timeout_per_batch=BATCH_REQ_TIMEOUT_DEFAULT,
+        timeout_total=None,
+        raise_on_fail=True,
+    ):
         """
         Get details for a batch of IP addresses at once.
 
@@ -158,31 +165,48 @@ async def getBatchDetails(self, ip_addresses, batch_size=None):
             else:
                 lookup_addresses.append(ip_address)
 
+        # all in cache - return early.
+        if len(lookup_addresses) == 0:
+            return result
+
+        # do start timer if necessary
+        if timeout_total is not None:
+            start_time = time.time()
+
         # loop over batch chunks and prepare coroutines for each.
+        url = API_URL + "/batch"
+        headers = handler_utils.get_headers(self.access_token)
+        headers["content-type"] = "application/json"
         reqs = []
-        for i in range(0, len(ip_addresses), batch_size):
-            chunk = ip_addresses[i : i + batch_size]
-
-            # all in cache - return early.
-            if len(lookup_addresses) == 0:
-                return result
+        for i in range(0, len(lookup_addresses), batch_size):
+            chunk = lookup_addresses[i : i + batch_size]
 
             # do http req
-            url = API_URL + "/batch"
-            headers = handler_utils.get_headers(self.access_token)
-            headers["content-type"] = "application/json"
             reqs.append(
                 self.httpsess.post(
-                    url, data=json.dumps(lookup_addresses), headers=headers
+                    url,
+                    data=json.dumps(chunk),
+                    headers=headers,
+                    timeout=timeout_per_batch,
                 )
             )
 
-        resps = await asyncio.gather(*reqs)
+        resps = await asyncio.wait_for(
+            asyncio.gather(*reqs, return_exceptions=raise_on_fail),
+            timeout_total
+        )
         for resp in resps:
             # gather data
-            if resp.status == 429:
-                raise RequestQuotaExceededError()
-            resp.raise_for_status()
+            try:
+                if resp.status == 429:
+                    raise RequestQuotaExceededError()
+                resp.raise_for_status()
+            except Exception as e:
+                if raise_on_fail:
+                    raise e
+                else:
+                    return result
+
             json_resp = await resp.json()
 
             # format & fill up cache