basic batch improvements.

chunking happens now; input list size limit doesnt exist.

Author: UmanShahzad

ipinfo/handler.py

@@ -83,15 +83,28 @@ def getDetails(self, ip_address=None):
 
         return Details(details)
 
-    def getBatchDetails(self, ip_addresses):
-        """Get details for a batch of IP addresses at once."""
+    def getBatchDetails(self, ip_addresses, batch_size=None):
+        """
+        Get details for a batch of IP addresses at once.
+
+        There is no specified limit to the number of IPs this function can
+        accept; it can handle as much as the user can fit in RAM (along with
+        all of the response data, which is at least a magnitude larger than the
+        input list).
+
+        The batch size can be adjusted with `batch_size` but is clipped to (and
+        also defaults to) `handler_utils.BATCH_MAX_SIZE`.
+        """
+        if batch_size == None:
+            batch_size = handler_utils.BATCH_MAX_SIZE
+
         result = {}
 
-        # Pre-populate with anything we've got in the cache, and keep around
+        # pre-populate with anything we've got in the cache, and keep around
         # the IPs not in the cache.
         lookup_addresses = []
         for ip_address in ip_addresses:
-            # If the supplied IP address uses the objects defined in the
+            # if the supplied IP address uses the objects defined in the
             # built-in module ipaddress extract the appropriate string notation
             # before formatting the URL.
             if isinstance(ip_address, IPv4Address) or isinstance(
@@ -104,28 +117,35 @@ def getBatchDetails(self, ip_addresses):
             else:
                 lookup_addresses.append(ip_address)
 
-        # Do the lookup
-        url = handler_utils.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, **self.request_options
-        )
-        if response.status_code == 429:
-            raise RequestQuotaExceededError()
-        response.raise_for_status()
-
-        # Fill up cache
-        json_response = response.json()
-        for ip_address, details in json_response.items():
-            self.cache[ip_address] = details
-
-        # Merge cached results with new lookup
-        result.update(json_response)
-
-        # Format every result
-        for detail in result.values():
-            if isinstance(detail, dict):
-                handler_utils.format_details(detail, self.countries)
+        # loop over batch chunks and do lookup for each.
+        for i in range(0, len(ip_addresses), batch_size):
+            chunk = ip_addresses[i : i + batch_size]
+
+            # lookup
+            url = handler_utils.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,
+                **self.request_options
+            )
+            if response.status_code == 429:
+                raise RequestQuotaExceededError()
+            response.raise_for_status()
+
+            # fill cache
+            json_response = response.json()
+            for ip_address, details in json_response.items():
+                self.cache[ip_address] = details
+
+            # merge cached results with new lookup
+            result.update(json_response)
+
+            # format all
+            for detail in result.values():
+                if isinstance(detail, dict):
+                    handler_utils.format_details(detail, self.countries)
 
         return result

ipinfo/handler_async.py

@@ -3,6 +3,7 @@
 """
 
 from ipaddress import IPv4Address, IPv6Address
+import asyncio
 import json
 import os
 import sys
@@ -112,10 +113,26 @@ async def getDetails(self, ip_address=None):
 
         return Details(details)
 
-    async def getBatchDetails(self, ip_addresses):
-        """Get details for a batch of IP addresses at once."""
+    async def getBatchDetails(self, ip_addresses, batch_size=None):
+        """
+        Get details for a batch of IP addresses at once.
+
+        There is no specified limit to the number of IPs this function can
+        accept; it can handle as much as the user can fit in RAM (along with
+        all of the response data, which is at least a magnitude larger than the
+        input list).
+
+        The batch size can be adjusted with `batch_size` but is clipped to (and
+        also defaults to) `handler_utils.BATCH_MAX_SIZE`.
+
+        The concurrency level is currently unadjustable; coroutines will be
+        created and consumed for all batches at once.
+        """
         self._ensure_aiohttp_ready()
 
+        if batch_size == None:
+            batch_size = handler_utils.BATCH_MAX_SIZE
+
         result = {}
 
         # Pre-populate with anything we've got in the cache, and keep around
@@ -135,32 +152,41 @@ async def getBatchDetails(self, ip_addresses):
             else:
                 lookup_addresses.append(ip_address)
 
-        # all in cache - return early.
-        if len(lookup_addresses) == 0:
-            return result
-
-        # do http req
-        url = handler_utils.API_URL + "/batch"
-        headers = handler_utils.get_headers(self.access_token)
-        headers["content-type"] = "application/json"
-        async with self.httpsess.post(
-            url,
-            data=json.dumps(lookup_addresses),
-            headers=headers
-        ) as resp:
+        # loop over batch chunks and prepare coroutines for each.
+        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
+
+            # do http req
+            url = handler_utils.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
+                )
+            )
+
+        resps = await asyncio.gather(*reqs)
+        for resp in resps:
+            # gather data
             if resp.status == 429:
                 raise RequestQuotaExceededError()
             resp.raise_for_status()
             json_resp = await resp.json()
 
-        # format & fill up cache
-        for ip_address, details in json_resp.items():
-            if isinstance(details, dict):
-                handler_utils.format_details(details, self.countries)
-                self.cache[ip_address] = details
+            # format & fill up cache
+            for ip_address, details in json_resp.items():
+                if isinstance(details, dict):
+                    handler_utils.format_details(details, self.countries)
+                    self.cache[ip_address] = details
 
-        # merge cached results with new lookup
-        result.update(json_resp)
+            # merge cached results with new lookup
+            result.update(json_resp)
 
         return result
 

ipinfo/handler_utils.py

@@ -8,15 +8,22 @@
 
 from .version import SDK_VERSION
 
+# Base URL to make requests against.
 API_URL = "https://ipinfo.io"
+
+# Used to transform incoming responses with country abbreviations into the full
+# expanded country name, e.g. "PK" -> "Pakistan".
 COUNTRY_FILE_DEFAULT = "countries.json"
 
+# The max amount of IPs allowed by the API per batch request.
+BATCH_MAX_SIZE = 1000
+
+
 def get_headers(access_token):
     """Build headers for request to IPinfo API."""
     headers = {
         "user-agent": "IPinfoClient/Python{version}/{sdk_version}".format(
-            version=sys.version_info[0],
-            sdk_version=SDK_VERSION
+            version=sys.version_info[0], sdk_version=SDK_VERSION
         ),
         "accept": "application/json",
     }
@@ -26,16 +33,16 @@ def get_headers(access_token):
 
     return headers
 
+
 def format_details(details, countries):
     """
     Format details given a countries object.
 
     The countries object can be retrieved from read_country_names.
     """
     details["country_name"] = countries.get(details.get("country"))
-    details["latitude"], details["longitude"] = read_coords(
-        details.get("loc")
-    )
+    details["latitude"], details["longitude"] = read_coords(details.get("loc"))
+
 
 def read_coords(location):
     """
@@ -50,6 +57,7 @@ def read_coords(location):
         lat, lon = coords[0], coords[1]
     return lat, lon
 
+
 def read_country_names(countries_file=None):
     """
     Read list of countries from specified country file or

ipinfo/version.py

@@ -1 +1 @@
-SDK_VERSION = '4.1.0'
+SDK_VERSION = "4.1.0"

tests/handler_async_test.py

@@ -29,9 +29,8 @@ async def test_headers():
     assert "authorization" in headers
 
 
-@pytest.mark.parametrize("n", range(5))
 @pytest.mark.asyncio
-async def test_get_details(n):
+async def test_get_details():
     token = os.environ.get("IPINFO_TOKEN", "")
     handler = AsyncHandler(token)
     details = await handler.getDetails("8.8.8.8")
@@ -86,27 +85,42 @@ async def test_get_details(n):
     await handler.deinit()
 
 
-@pytest.mark.parametrize("n", range(5))
-@pytest.mark.asyncio
-async def test_get_batch_details(n):
+#############
+# BATCH TESTS
+#############
+
+_batch_ip_addrs = ["1.1.1.1", "8.8.8.8", "9.9.9.9"]
+
+
+def _prepare_batch_test():
+    """Helper for preparing batch test cases."""
     token = os.environ.get("IPINFO_TOKEN", "")
     if not token:
         pytest.skip("token required for batch tests")
     handler = AsyncHandler(token)
-    ips = ["1.1.1.1", "8.8.8.8", "9.9.9.9"]
-    details = await handler.getBatchDetails(ips)
+    return handler, token, _batch_ip_addrs
 
+
+def _check_batch_details(ips, details, token):
+    """Helper for batch tests."""
     for ip in ips:
         assert ip in details
         d = details[ip]
         assert d["ip"] == ip
-        assert d["country"] == "US"
-        assert d["country_name"] == "United States"
+        assert "country" in d
+        assert "country_name" in d
         if token:
             assert "asn" in d
             assert "company" in d
             assert "privacy" in d
             assert "abuse" in d
             assert "domains" in d
 
+
+@pytest.mark.parametrize("batch_size", [None, 2, 3])
+@pytest.mark.asyncio
+async def test_get_batch_details(batch_size):
+    handler, token, ips = _prepare_batch_test()
+    details = await handler.getBatchDetails(ips, batch_size=batch_size)
+    _check_batch_details(ips, details, token)
     await handler.deinit()

tests/handler_test.py

@@ -27,8 +27,7 @@ def test_headers():
     assert "authorization" in headers
 
 
-@pytest.mark.parametrize("n", range(5))
-def test_get_details(n):
+def test_get_details():
     token = os.environ.get("IPINFO_TOKEN", "")
     handler = Handler(token)
     details = handler.getDetails("8.8.8.8")
@@ -81,24 +80,40 @@ def test_get_details(n):
         assert len(domains["domains"]) == 5
 
 
-@pytest.mark.parametrize("n", range(5))
-def test_get_batch_details(n):
+#############
+# BATCH TESTS
+#############
+
+_batch_ip_addrs = ["1.1.1.1", "8.8.8.8", "9.9.9.9"]
+
+
+def _prepare_batch_test():
+    """Helper for preparing batch test cases."""
     token = os.environ.get("IPINFO_TOKEN", "")
     if not token:
         pytest.skip("token required for batch tests")
     handler = Handler(token)
-    ips = ["1.1.1.1", "8.8.8.8", "9.9.9.9"]
-    details = handler.getBatchDetails(ips)
+    return handler, token, _batch_ip_addrs
 
+
+def _check_batch_details(ips, details, token):
+    """Helper for batch tests."""
     for ip in ips:
         assert ip in details
         d = details[ip]
         assert d["ip"] == ip
-        assert d["country"] == "US"
-        assert d["country_name"] == "United States"
+        assert "country" in d
+        assert "country_name" in d
         if token:
             assert "asn" in d
             assert "company" in d
             assert "privacy" in d
             assert "abuse" in d
             assert "domains" in d
+
+
+@pytest.mark.parametrize("batch_size", [None, 2, 3])
+def test_get_batch_details(batch_size):
+    handler, token, ips = _prepare_batch_test()
+    details = handler.getBatchDetails(ips, batch_size=batch_size)
+    _check_batch_details(ips, details, token)