adding error and cache

Author: Samuel-Gill

ipinfo/handler.py

@@ -306,23 +306,89 @@ def getMap(self, ips):
         response.raise_for_status()
         return response.json()["reportUrl"]
 
-    def getIterativeBatchDetails(self, ip_addresses, batch_size=None):
+    def getIterativeBatchDetails(
+        self,
+        ip_addresses,
+        batch_size=None,
+        raise_on_fail=True,
+    ):
         if batch_size == None:
             batch_size = BATCH_MAX_SIZE
 
+        results = {}
+
+        # 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
+            # built-in module ipaddress extract the appropriate string notation
+            # before formatting the URL.
+            if isinstance(ip_address, IPv4Address) or isinstance(
+                ip_address, IPv6Address
+            ):
+                ip_address = ip_address.exploded
+
+            # check if bogon.
+            if ip_address and is_bogon(ip_address):
+                details = {}
+                details["ip"] = ip_address
+                details["bogon"] = True
+                yield Details(details)
+
+            # check cache first.
+            try:
+                cached_ipaddr = self.cache[cache_key(ip_address)]
+                results[ip_address] = cached_ipaddr
+            except KeyError:
+                lookup_addresses.append(ip_address)
+
+        # all in cache - return early.
+        if len(lookup_addresses) == 0:
+            for ip_address, details in results.items():
+                yield ip_address, details
+
+        # loop over batch chunks and do lookup for each.
         url = API_URL + "/batch"
         headers = handler_utils.get_headers(self.access_token, self.headers)
         headers["content-type"] = "application/json"
+        for i in range(0, len(lookup_addresses), batch_size):
+            batch = lookup_addresses[i : i + batch_size]
 
-        # Split the IP addresses into batches
-        batches = [
-            ip_addresses[i : i + batch_size]
-            for i in range(0, len(ip_addresses), batch_size)
-        ]
+            # lookup.
+            try:
+                response = requests.post(url, json=batch, headers=headers)
+            except Exception as e:
+                return handler_utils.return_or_fail(raise_on_fail, e, result)
 
-        for batch in batches:
-            response = requests.post(url, json=batch, headers=headers)
-            json_response = response.json()
+            # fail on bad status codes
+            try:
+                if response.status_code == 429:
+                    raise RequestQuotaExceededError()
+                response.raise_for_status()
+            except Exception as e:
+                return handler_utils.return_or_fail(raise_on_fail, e, result)
 
+            # fill cache
+            json_response = response.json()
+            print(f"JSON: {json_response}")
             for ip_address, details in json_response.items():
+                self.cache[cache_key(ip_address)] = details
+
+            # merge cached results with new lookup
+            results.update(json_response)
+
+            # format all
+            for detail in results.values():
+                if isinstance(detail, dict):
+                    handler_utils.format_details(
+                        detail,
+                        self.countries,
+                        self.eu_countries,
+                        self.countries_flags,
+                        self.countries_currencies,
+                        self.continents,
+                    )
+
+            for ip_address, details in results.items():
                 yield ip_address, details

ipinfo/handler_utils.py

@@ -79,7 +79,9 @@ def format_details(
     """
     details["country_name"] = countries.get(details.get("country"))
     details["isEU"] = details.get("country") in eu_countries
-    details["country_flag_url"] = COUNTRY_FLAGS_URL + details.get("country") + ".svg"
+    details["country_flag_url"] = (
+        COUNTRY_FLAGS_URL + details.get("country") + ".svg"
+    )
     details["country_flag"] = copy.deepcopy(
         countries_flags.get(details.get("country"))
     )

tests/handler_async_test.py

@@ -48,7 +48,10 @@ async def test_get_details():
     assert country_flag["emoji"] == "🇺🇸"
     assert country_flag["unicode"] == "U+1F1FA U+1F1F8"
     country_flag_url = details.country_flag_url
-    assert country_flag_url == "https://cdn.ipinfo.io/static/images/countries-flags/US.svg"
+    assert (
+        country_flag_url
+        == "https://cdn.ipinfo.io/static/images/countries-flags/US.svg"
+    )
     country_currency = details.country_currency
     assert country_currency["code"] == "USD"
     assert country_currency["symbol"] == "$"

tests/handler_test.py

@@ -45,7 +45,10 @@ def test_get_details():
     assert country_flag["emoji"] == "🇺🇸"
     assert country_flag["unicode"] == "U+1F1FA U+1F1F8"
     country_flag_url = details.country_flag_url
-    assert country_flag_url == "https://cdn.ipinfo.io/static/images/countries-flags/US.svg"
+    assert (
+        country_flag_url
+        == "https://cdn.ipinfo.io/static/images/countries-flags/US.svg"
+    )
     country_currency = details.country_currency
     assert country_currency["code"] == "USD"
     assert country_currency["symbol"] == "$"
@@ -189,3 +192,11 @@ def test_bogon_details():
     details = handler.getDetails("127.0.0.1")
     assert isinstance(details, Details)
     assert details.all == {"bogon": True, "ip": "127.0.0.1"}
+
+
+def test_iterative_bogon_details():
+    token = os.environ.get("IPINFO_TOKEN", "")
+    handler = Handler(token)
+    details = next(handler.getIterativeBatchDetails(["127.0.0.1"]))
+    assert isinstance(details, Details)
+    assert details.all == {"bogon": True, "ip": "127.0.0.1"}