From 427a9a2ecea16cfaaea04a4bcedae6017d70be8a Mon Sep 17 00:00:00 2001
From: RamanjaneyuluIdavalapati <ramanjaneyulu@deepcompute.com>
Date: Tue, 14 May 2019 11:16:09 +0000
Subject: [PATCH] Supported setting response headers

---
 .travis.yml       |  4 ++--
 README.md         |  4 ++--
 kwikapi/api.py    | 22 ++++++++++++++++++----
 kwikapi/client.py |  9 ++++-----
 setup.py          |  8 ++++----
 5 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/.travis.yml b/.travis.yml
index 1205b5e..ce11ea5 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -16,8 +16,8 @@ deploy:
     - LICENSE
     - kwikapi/api.py
     - kwikapi/__init__.py
-  name: kwikapi-0.5.6
-  tag_name: 0.5.6
+  name: kwikapi-0.5.7
+  tag_name: 0.5.7
   on:
     repo: deep-compute/kwikapi
 - provider: pypi
diff --git a/README.md b/README.md
index 612cd5c..20a75b7 100644
--- a/README.md
+++ b/README.md
@@ -510,12 +510,12 @@ $ wget "http://localhost:8888/api/v1/add" --header="X-KwikAPI-Protocol: numpy" -
 $ wget "http://localhost:8888/api/v1/add?a=10&b=20" --header="X-KwikAPI-Protocol: raw"
 ```
 
-#### We can also change outgoing protocol
+#### We can also change outgoing/response protocol
 ex:
 ```python
 class Calc(object):
     def add(self, req: Request, a: int, b: int) -> int:
-        req.headers['X-KwikAPI-Protocol'] = 'numpy'
+        req.response.headers['X-KwikAPI-Protocol'] = 'numpy'
         return a + b
 ```
 
diff --git a/kwikapi/api.py b/kwikapi/api.py
index 77ea100..45fbe8b 100644
--- a/kwikapi/api.py
+++ b/kwikapi/api.py
@@ -125,6 +125,10 @@ def flush(self):
     def close(self):
         pass
 
+    @abc.abstractproperty
+    def headers(self):
+        pass
+
 class MockRequest(BaseRequest):
 
     def __init__(self, **kwargs):
@@ -153,7 +157,7 @@ def headers(self):
 class MockResponse(BaseResponse):
     def __init__(self):
         super().__init__()
-        self.headers = CaseInsensitiveDict()
+        self._headers = CaseInsensitiveDict()
         self.raw_response = None
 
     def write(self, data, protocol, stream=False):
@@ -169,6 +173,10 @@ def flush(self):
     def close(self):
         pass
 
+    @property
+    def headers(self):
+        return self._headers
+
 class API(object):
     """
     A collection of APIFragments
@@ -460,10 +468,16 @@ def _resolve_call_info(self, request):
 
         return r
 
-    def _find_request_protocol(self, request):
-        protocol = request.headers.get(PROTOCOL_HEADER, self.default_protocol)
+    def _find_request_protocol(self, r):
+        protocol = r.headers.get(PROTOCOL_HEADER, self.default_protocol)
         return self.PROTOCOLS[protocol]
 
+    def _find_response_protocol(self, r):
+        protocol = r.response.headers.get(PROTOCOL_HEADER, None)
+        if protocol:
+            return self.PROTOCOLS[protocol]
+        return self._find_request_protocol(r)
+
     def _handle_exception(self, req, e):
         message_value = e.message if hasattr(e, 'message') else str(e)
         code_value = e.code if hasattr(e, 'code') else self.DEFAULT_ERROR_CODE
@@ -520,7 +534,7 @@ def handle_request(self, request):
                 self._invoke_pre_call_hook(request)
                 result = request.fn(**request.fn_params)
 
-                protocol = self._find_request_protocol(request)
+                protocol = self._find_response_protocol(request)
 
                 self._invoke_post_call_hook(request, result=result)
 
diff --git a/kwikapi/client.py b/kwikapi/client.py
index 8317e6b..0841db8 100644
--- a/kwikapi/client.py
+++ b/kwikapi/client.py
@@ -111,19 +111,18 @@ def _make_request(self, url, post_body, headers):
         req = urllib.request.Request(url, data=post_body, headers=headers)
         res = urllib.request.urlopen(req)
 
+        proto = PROTOCOLS[res.headers.get('X-KwikAPI-Protocol', self._protocol)]
+
         if self._stream:
-            proto = PROTOCOLS[self._protocol]
             res = proto.deserialize_stream(res)
             res = Client._extract_stream_response(res, self._raise_exception)
         else:
-            res = self._deserialize_response(res.read(), self._protocol,
-                    self._raise_exception)
+            res = self._deserialize_response(res.read(), proto, self._raise_exception)
 
         return res
 
     @staticmethod
-    def _deserialize_response(data, protocol, raise_exception=True):
-        proto = PROTOCOLS[protocol]
+    def _deserialize_response(data, proto, raise_exception=True):
         r = proto.deserialize(data)
         return Client._extract_response(r, raise_exception)
 
diff --git a/setup.py b/setup.py
index f1f00d9..486c40b 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
 from setuptools import setup, find_packages
 
-version = '0.5.6'
+version = '0.5.7'
 setup(
     name="kwikapi",
     version=version,
@@ -21,9 +21,9 @@
     'requests>=2.18.4',
     ],
     extras_require={
-        'django': ['kwikapi-django==0.2.5'],
-        'tornado': ['kwikapi-tornado==0.3.4'],
-        'all': ['kwikapi-django==0.2.5', 'kwikapi-tornado==0.3.4']
+        'django': ['kwikapi-django==0.2.6'],
+        'tornado': ['kwikapi-tornado==0.3.7'],
+        'all': ['kwikapi-django==0.2.6', 'kwikapi-tornado==0.3.7']
     },
     classifiers=[
         'Environment :: Web Environment',