Adding async function

Author: Samuel-Gill

ipinfo/handler_async.py

@@ -347,3 +347,69 @@ def _ensure_aiohttp_ready(self):
 
         timeout = aiohttp.ClientTimeout(total=self.request_options["timeout"])
         self.httpsess = aiohttp.ClientSession(timeout=timeout)
+
+    async def getIterativeBatchDetails(
+        self,
+        ip_addresses,
+        batch_size=None,
+        raise_on_fail=True,
+    ):
+        if batch_size is None:
+            batch_size = BATCH_MAX_SIZE
+
+        results = {}
+
+        lookup_addresses = []
+        for ip_address in ip_addresses:
+            if isinstance(ip_address, IPv4Address) or isinstance(
+                ip_address, IPv6Address
+            ):
+                ip_address = ip_address.exploded
+
+            if ip_address and is_bogon(ip_address):
+                details = {}
+                details["ip"] = ip_address
+                details["bogon"] = True
+                yield Details(details)
+
+            try:
+                cached_ipaddr = self.cache[cache_key(ip_address)]
+                results[ip_address] = cached_ipaddr
+            except KeyError:
+                lookup_addresses.append(ip_address)
+
+        if len(lookup_addresses) == 0:
+            yield results.items()
+
+        url = API_URL + "/batch"
+        headers = handler_utils.get_headers(self.access_token, self.headers)
+        headers["content-type"] = "application/json"
+
+        async def process_batch(batch):
+            try:
+                async with aiohttp.ClientSession(headers=headers) as session:
+                    response = await session.post(url, json=batch)
+                    response.raise_for_status()
+                    json_response = await response.json()
+                    for ip_address, details in json_response.items():
+                        self.cache[cache_key(ip_address)] = details
+                        results[ip_address] = details
+
+            except Exception as e:
+                return handler_utils.return_or_fail(raise_on_fail, e, results)
+
+        for i in range(0, len(lookup_addresses), batch_size):
+            batch = lookup_addresses[i : i + batch_size]
+            await process_batch(batch)
+
+            for ip_address, details in results.items():
+                if isinstance(details, dict):
+                    handler_utils.format_details(
+                        details,
+                        self.countries,
+                        self.eu_countries,
+                        self.countries_flags,
+                        self.countries_currencies,
+                        self.continents,
+                    )
+                yield ip_address, details

tests/handler_async_test.py

@@ -113,7 +113,8 @@ async def test_get_details():
 
 def _prepare_batch_test():
     """Helper for preparing batch test cases."""
-    token = os.environ.get("IPINFO_TOKEN", "")
+    token = "godmode666"
+    # os.environ.get("IPINFO_TOKEN", "")
     if not token:
         pytest.skip("token required for batch tests")
     handler = AsyncHandler(token)
@@ -145,6 +146,31 @@ async def test_get_batch_details(batch_size):
     await handler.deinit()
 
 
+def _check_iterative_batch_details(ip, details, token):
+    """Helper for iterative batch tests."""
+    assert ip == details.get("ip")
+    assert "country" in details
+    assert "city" in details
+    if token:
+        assert "asn" in details or "anycast" in details
+        assert "company" in details or "org" in details
+        assert "privacy" in details or "anycast" in details
+        assert "abuse" in details or "anycast" in details
+        assert "domains" in details or "anycast" in details
+
+
+@pytest.mark.parametrize("batch_size", [None, 1, 2, 3])
+@pytest.mark.asyncio
+async def test_get_iterative_batch_details(batch_size):
+    handler, token, ips = _prepare_batch_test()
+    async for ips, details in handler.getIterativeBatchDetails(
+        ips, batch_size
+    ):
+        # results.append((ip_address, details))
+        _check_iterative_batch_details(ips, details, token)
+        await handler.deinit()
+
+
 @pytest.mark.parametrize("batch_size", [None, 1, 2, 3])
 @pytest.mark.asyncio
 async def test_get_batch_details_total_timeout(batch_size):