Merge pull request #49 from SamHausmann/RCB-412

added custom headers capability

Author: cp2boston

examples/language.py

@@ -17,8 +17,8 @@ def run(key, altUrl='https://api.rosette.com/rest/v1/'):
 
     language_data = "Por favor Señorita, says the man."
     params = DocumentParameters()
-
     params["content"] = language_data
+    api.setCustomHeaders("X-RosetteAPI-App", "python-app")
     return api.language(params)
 
 

rosette/api.py

@@ -26,6 +26,7 @@
 import os
 from socket import gaierror
 import requests
+import re
 
 _BINDING_VERSION = '1.1.1'
 _GZIP_BYTEARRAY = bytearray([0x1F, 0x8b, 0x08])
@@ -366,6 +367,17 @@ def info(self):
         identifying data."""
         url = self.service_url + "info"
         headers = {'Accept': 'application/json', 'X-RosetteAPI-Binding': 'python', 'X-RosetteAPI-Binding-Version': _BINDING_VERSION}
+
+        customHeaders = self.api.getCustomHeaders()
+        pattern = re.compile('^X-RosetteAPI-')
+        if customHeaders is not None:
+            for key in customHeaders.keys():
+                if pattern.match(key) is not None:
+                    headers[key] = customHeaders[key]
+                else:
+                    raise RosetteException("badHeader", "Custom header name must begin with \"X-RosetteAPI-\"", key)
+            self.api.clearCustomHeaders()
+
         if self.debug:
             headers['X-RosetteAPI-Devel'] = 'true'
         self.logger.info('info: ' + url)
@@ -382,6 +394,17 @@ def ping(self):
 
         url = self.service_url + 'ping'
         headers = {'Accept': 'application/json', 'X-RosetteAPI-Binding': 'python', 'X-RosetteAPI-Binding-Version': _BINDING_VERSION}
+
+        customHeaders = self.api.getCustomHeaders()
+        pattern = re.compile('^X-RosetteAPI-')
+        if customHeaders is not None:
+            for key in customHeaders.keys():
+                if pattern.match(key) is not None:
+                    headers[key] = customHeaders[key]
+                else:
+                    raise RosetteException("badHeader", "Custom header name must begin with \"X-RosetteAPI-\"", key)
+            self.api.clearCustomHeaders()
+
         if self.debug:
             headers['X-RosetteAPI-Devel'] = 'true'
         self.logger.info('Ping: ' + url)
@@ -426,9 +449,21 @@ def call(self, parameters):
         params_to_serialize = parameters.serialize(self.api.options)
         headers = {}
         if self.user_key is not None:
+
+            customHeaders = self.api.getCustomHeaders()
+            pattern = re.compile('^X-RosetteAPI-')
+            if customHeaders is not None:
+                for key in customHeaders.keys():
+                    if pattern.match(key) is not None:
+                        headers[key] = customHeaders[key]
+                    else:
+                        raise RosetteException("badHeader", "Custom header name must begin with \"X-RosetteAPI-\"", key)
+                self.api.clearCustomHeaders()
+
             headers["X-RosetteAPI-Key"] = self.user_key
             headers["X-RosetteAPI-Binding"] = "python"
             headers["X-RosetteAPI-Binding-Version"] = _BINDING_VERSION
+
         if self.useMultipart:
             params = dict(
                 (key,
@@ -502,6 +537,7 @@ def __init__(
         self.connection_refresh_duration = refresh_duration
         self.http_connection = None
         self.options = {}
+        self.customHeaders = {}
 
     def _connect(self, parsedUrl):
         """ Simple connection method
@@ -638,6 +674,30 @@ def clearOptions(self):
         """
         self.options.clear()
 
+    def setCustomHeaders(self, name, value):
+        """
+        Sets custom headers
+
+        @param headers: array of custom headers to be set
+        """
+        if value is None:
+            self.customHeaders.pop(name, None)
+        else:
+            self.customHeaders[name] = value
+
+    def getCustomHeaders(self):
+        """
+        Get custom headers
+        """
+        return self.customHeaders
+
+    def clearCustomHeaders(self):
+        """
+        Clears custom headers
+        """
+
+        self.customHeaders.clear()
+
     def ping(self):
         """
         Create a ping L{EndpointCaller} for the server and ping it.

tests/test_rosette_api.py

@@ -82,10 +82,35 @@ def test_option_clear_single_option(api):
     api.setOption('test', None)
     assert api.getOption('test') is None
 
+# Test the custom header set/get/clear
+
+
+def test_custom_header_get_set_clear(api):
+    key = 'X-RosetteAPI-Test'
+    value = 'foo'
+    api.setCustomHeaders(key, value)
+    assert value == api.getCustomHeaders()[key]
+
+    api.clearCustomHeaders()
+    assert len(api.getCustomHeaders()) is 0
+
+# Test for invalid header name
+
+
+def test_invalid_header(api):
+    key = 'test'
+    value = 'foo'
+    api.setCustomHeaders(key, value)
+
+    with pytest.raises(RosetteException) as e_rosette:
+        result = api.info()
+
+    assert e_rosette.value.status == 'badHeader'
 
 # Test that pinging the API is working properly
 # @httpretty.activate
 
+
 def test_ping(api, json_response):
     httpretty.enable()
     httpretty.register_uri(httpretty.GET, "https://api.rosette.com/rest/v1/ping",