diff --git a/README.md b/README.md index ec83435..a5ce03c 100755 --- a/README.md +++ b/README.md @@ -148,6 +148,49 @@ result = lambda_handler(event=valid_input_event) assert result == {"body": '{"foo": [1.0, 2.2, 3.0]}', "statusCode": 200, "headers":{}} ``` +### Headers and MultiValueHeaders + +You can return headers as part of your response, there is two types of headers: + +* **headers**: a dictionary +* **multiValueHeaders**: a dictionary with multiple values for each key + +You can populate the headers in the following ways: + +1. return tuple with *normal header* + +```python +return {'some':'json'}, 200, {'header':'value'} +``` + +2. return tuple with *normal header* **and** *multiValueHeaders* + +```json +return {'some':'json'}, 200, {'header':'value'}, {'multi_header': ["foo", "bar"]} +``` + +3. return tuple with *multiValueHeaders* + +*(you need to still populate the headers as an empty dict)* + +```json +return {'some':'json'}, 200, {}, {'multi_header': ["foo", "bar"]} +``` + +4. return [`Response`](https://github.com/trustpilot/python-lambdarest/blob/cd50bb4e1da4f720ef94534ccfd4989f398a9d5d/lambdarest/__init__.py#L17) object + +```python +from lambdarest import Response +# with headers +Response({'some':'json'}, 200, headers={'header':'value'}) + +# with multiValueHeaders +Response({'some':'json'}, 200, multiValueHeaders={'multi_header': ["foo", "bar"]}) + +# with headers and multiValueHeaders +Response({'some':'json'}, 200, headers={'header':'value'}, multiValueHeaders={'multi_header': ["foo", "bar"]}) +``` + ### Routing You can also specify which path to react on for individual handlers using the `path` param: diff --git a/lambdarest/__init__.py b/lambdarest/__init__.py index d8cd80f..280b126 100755 --- a/lambdarest/__init__.py +++ b/lambdarest/__init__.py @@ -22,10 +22,13 @@ class Response(object): if no headers are specified, empty dict is returned """ - def __init__(self, body=None, status_code=None, headers=None): + def __init__( + self, body=None, status_code=None, headers=None, multiValueHeaders=None + ): self.body = body self.status_code = status_code self.headers = headers + self.multiValueHeaders = multiValueHeaders self.status_code_description = None self.isBase64_encoded = False @@ -42,8 +45,12 @@ def to_json(self, encoder=json.JSONEncoder, application_load_balancer=False): if do_json_dumps else self.body, "statusCode": status_code, - "headers": self.headers or {}, } + ## handle multiValueHeaders if defined, default to headers + if self.multiValueHeaders == None: + response["headers"] = self.headers or {} + else: + response["multiValueHeaders"] = self.multiValueHeaders # if body is None, remove the key if response.get("body") == None: response.pop("body") @@ -227,7 +234,7 @@ def inner_lambda_handler(event, context=None): response = func(event, **kwargs) if not isinstance(response, Response): # Set defaults - status_code = headers = None + status_code = headers = multiValueHeaders = None if isinstance(response, tuple): response_len = len(response) @@ -235,21 +242,24 @@ def inner_lambda_handler(event, context=None): raise ValueError("Response tuple has more than 3 items") # Unpack the tuple, missing items will be defaulted - body, status_code, headers = response + (None,) * ( - 3 - response_len - ) + body, status_code, headers, multiValueHeaders = response + ( + None, + ) * (4 - response_len) elif isinstance(response, dict) and all( - key in ["body", "statusCode", "headers"] + key in ["body", "statusCode", "headers", "multiValueHeaders"] for key in response.keys() ): body = response.get("body") status_code = response.get("statusCode") or status_code headers = response.get("headers") or headers + multiValueHeaders = ( + response.get("multiValueHeaders") or multiValueHeaders + ) else: # if response is string, int, etc. body = response - response = Response(body, status_code, headers) + response = Response(body, status_code, headers, multiValueHeaders) return response.to_json( encoder=json_encoder, application_load_balancer=application_load_balancer, diff --git a/tests/test_lambdarest.py b/tests/test_lambdarest.py index 96f53ee..4f56877 100755 --- a/tests/test_lambdarest.py +++ b/tests/test_lambdarest.py @@ -485,6 +485,25 @@ def test_that_standard_dict_responses_without_a_body_are_returned_as_is(self): result, {"headers": {"Location": "https://example.com"}, "statusCode": 302} ) + def test_that_standard_dict_responses_without_a_body_are_returned_as_is_multiHeader( + self, + ): + post_mock = mock.Mock( + return_value={ + "statusCode": 302, + "multiValueHeaders": {"Location": ["https://example.com"]}, + } + ) + self.lambda_handler.handle("post")(post_mock) # decorate mock + result = self.lambda_handler(self.event, self.context) + self.assertEqual( + result, + { + "multiValueHeaders": {"Location": ["https://example.com"]}, + "statusCode": 302, + }, + ) + def test_that_dict_responses_that_happen_to_have_a_body_key_retain_previous_behavior( self, ):