Merge pull request #79 from ipinfo/custom_headers

allow custom headers

Author: UmanShahzad

README.md

@@ -230,6 +230,14 @@ Request behavior can be modified by setting the `request_options` keyword argume
 >>> handler = ipinfo.getHandler(request_options={'timeout': 4})
 ```
 
+### Custom Headers
+
+You can add custom headers or modify default headers by setting the `headers` keyword argument when initializing the handler. `headers` is a dictionary of `{'header': 'value'}` format.
+
+```python
+>>> handler = ipinfo.getHandler(headers={'user-agent': 'My Custom User-agent', 'custom_header': 'yes'})
+```
+
 ### Internationalization
 
 When looking up an IP address, the response object includes a `details.country_name`, `details.isEU`, `details.country_flag` and `details.country_currency` attributes which includes the country based on American English. It is possible to return the country name in other languages by setting the `countries_file`, remove or add EU countries by setting the keyword argument `eu_countries_file`, change the country flag emoji or unicode by setting the keyword argument `countries_flags_file` or change country's currency code or currency symbol by setting the `countries_currencies` when creating the `IPinfo` object. Moreover the response object includes a `details.continent` which includes continent code and name of IP. The default file can be changed by setting the `continent_file` while creating the `IPinfo` object.

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,15 +44,18 @@
 BATCH_REQ_TIMEOUT_DEFAULT = 5
 
 
-def get_headers(access_token):
+def get_headers(access_token, custom_headers):
     """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
+       "user-agent": "IPinfoClient/Python{version}/{sdk_version}".format(
+        version=sys.version_info[0], sdk_version=SDK_VERSION
         ),
         "accept": "application/json",
     }
 
+    if custom_headers:
+        headers = {** headers, ** custom_headers}
+
     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={"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={"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():