allowed custom headers

Author: rm-Umar

ipinfo/handler.py

@@ -95,6 +95,9 @@ def __init__(self, access_token=None, **kwargs):
                 cache_options["ttl"] = CACHE_TTL
             self.cache = DefaultCache(**cache_options)
 
+        # setup custom headers
+        self.headers = kwargs.get("headers", None)
+
     def getDetails(self, ip_address=None, timeout=None):
         """
         Get details for specified IP address as a Details object.
@@ -133,7 +136,7 @@ def getDetails(self, ip_address=None, timeout=None):
         url = API_URL
         if ip_address:
             url += "/" + ip_address
-        headers = handler_utils.get_headers(self.access_token)
+        headers = handler_utils.get_headers(self.access_token, self.headers)
         response = requests.get(url, headers=headers, **req_opts)
         if response.status_code == 429:
             raise RequestQuotaExceededError()
@@ -226,7 +229,7 @@ def getBatchDetails(
 
         # loop over batch chunks and do lookup for each.
         url = API_URL + "/batch"
-        headers = handler_utils.get_headers(self.access_token)
+        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):
             # quit if total timeout is reached.
@@ -295,7 +298,7 @@ def getMap(self, ips):
 
         req_opts = {**self.request_options}
         url = f"{API_URL}/map?cli=1"
-        headers = handler_utils.get_headers(None)
+        headers = handler_utils.get_headers(None, self.headers)
         headers["content-type"] = "application/json"
         response = requests.post(
             url, json=ip_strs, headers=headers, **req_opts

ipinfo/handler_async.py

@@ -98,7 +98,10 @@ def __init__(self, access_token=None, **kwargs):
             if "ttl" not in cache_options:
                 cache_options["ttl"] = CACHE_TTL
             self.cache = DefaultCache(**cache_options)
-
+        
+                # setup custom headers
+        self.headers = kwargs.get("headers", None)
+        
     async def init(self):
         """
         Initializes internal aiohttp connection pool.
@@ -153,7 +156,7 @@ async def getDetails(self, ip_address=None, timeout=None):
         url = API_URL
         if ip_address:
             url += "/" + ip_address
-        headers = handler_utils.get_headers(self.access_token)
+        headers = handler_utils.get_headers(self.access_token, self.headers)
         req_opts = {}
         if timeout is not None:
             req_opts["timeout"] = timeout
@@ -251,7 +254,7 @@ async def getBatchDetails(
 
         # loop over batch chunks and prepare coroutines for each.
         url = API_URL + "/batch"
-        headers = handler_utils.get_headers(self.access_token)
+        headers = handler_utils.get_headers(self.access_token, self.headers)
         headers["content-type"] = "application/json"
 
         # prepare coroutines that will make reqs and update results.

ipinfo/handler_utils.py

@@ -44,14 +44,19 @@
 BATCH_REQ_TIMEOUT_DEFAULT = 5
 
 
-def get_headers(access_token):
-    """Build headers for request to IPinfo API."""
-    headers = {
+def get_headers(access_token, custom_headers):
+    headers = {}
+
+    if custom_headers:
+        headers = custom_headers
+    else:
+        """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
         ),
         "accept": "application/json",
-    }
+        }
 
     if access_token:
         headers["authorization"] = "Bearer {}".format(access_token)

tests/handler_async_test.py

@@ -21,13 +21,14 @@ async def test_init():
 @pytest.mark.asyncio
 async def test_headers():
     token = "mytesttoken"
-    handler = AsyncHandler(token)
-    headers = handler_utils.get_headers(token)
+    handler = AsyncHandler(token, headers={"user-agent": "test-agent", "accept": "application/json", "custom_field": "yes"})
+    headers = handler_utils.get_headers(token, handler.headers)
     await handler.deinit()
 
     assert "user-agent" in headers
     assert "accept" in headers
     assert "authorization" in headers
+    assert "custom_field" in headers
 
 
 @pytest.mark.asyncio

tests/handler_test.py

@@ -20,12 +20,13 @@ def test_init():
 
 def test_headers():
     token = "mytesttoken"
-    handler = Handler(token)
-    headers = handler_utils.get_headers(token)
+    handler = Handler(token, headers={"user-agent": "test-agent", "accept": "application/json", "custom_field": "yes"})
+    headers = handler_utils.get_headers(token, handler.headers)
 
     assert "user-agent" in headers
     assert "accept" in headers
     assert "authorization" in headers
+    assert "custom_field" in headers
 
 
 def test_get_details():