Merge pull request #3 from tofarr/web_page_updates

Web page updates

servey/servey_web_page/web_page_action_endpoint.py

@@ -1,12 +1,13 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from http.client import HTTPException
-from typing import Optional, Any
+from typing import Optional, Any, Dict
 
 from marshy.types import ExternalItemType
 from starlette.responses import Response, RedirectResponse
 
 from servey.servey_starlette.action_endpoint.action_endpoint import ActionEndpoint
 from servey.servey_web_page.redirect import Redirect
+from servey.servey_web_page.web_page_response import WebPageResponse
 from servey.servey_web_page.web_page_trigger import get_environment
 
 
@@ -17,7 +18,9 @@ class WebPageActionEndpoint(ActionEndpoint):
     """
 
     template_name: Optional[str] = None
-    content_type: str = "text/html"
+    response_headers: Dict[str, str] = field(
+        default_factory=lambda: {"Content-Type": "text/html"}
+    )
 
     def __post_init__(self):
         if not self.template_name:
@@ -26,15 +29,20 @@ def __post_init__(self):
     def render_response(self, result: Any):
         if isinstance(result, Redirect):
             return RedirectResponse(result.url, result.status_code)
+        if not isinstance(result, WebPageResponse):
+            result = WebPageResponse(result, headers=self.response_headers)
         result_content = (
-            self.result_marshaller.dump(result) if self.result_marshaller else None
+            self.result_marshaller.dump(result.model)
+            if self.result_marshaller
+            else None
         )
         if self.result_schema:
             error = next(self.result_schema.iter_errors(result_content), None)
             if error:
                 raise HTTPException(500, str(error))
         body = self.template.render(model=result_content)
-        return Response(content=body, media_type=self.content_type)
+
+        return Response(content=body, headers=result.headers)
 
     @property
     def template(self):

servey/servey_web_page/web_page_action_endpoint_factory.py

@@ -45,6 +45,7 @@ def create(
         result_type = inspect.signature(action.fn).return_annotation
         if result_type == inspect.Signature.empty:
             result_type = None
+        content_type = mimetypes.guess_type(action.name or "")[0] or "text/html"
         endpoint = WebPageActionEndpoint(
             action=action,
             path=path,
@@ -62,6 +63,6 @@ def create(
             if self.validate_output and result_type
             else None,
             template_name=trigger.template_name,
-            content_type=mimetypes.guess_type(action.name or "")[0] or "text/html",
+            response_headers={"Content-Type": content_type},
         )
         return endpoint

servey/servey_web_page/web_page_event_handler.py

@@ -2,7 +2,7 @@
 import inspect
 import mimetypes
 from dataclasses import field, dataclass
-from typing import Optional, Awaitable
+from typing import Optional, Awaitable, Dict
 
 from marshy import get_default_context
 from marshy.marshaller_context import MarshallerContext
@@ -20,13 +20,16 @@
 from servey.servey_aws.event_handler.event_handler import separate_auth_kwarg
 from servey.servey_aws.event_handler.event_handler_abc import EventHandlerFactoryABC
 from servey.servey_web_page.redirect import Redirect
+from servey.servey_web_page.web_page_response import WebPageResponse
 from servey.servey_web_page.web_page_trigger import WebPageTrigger, get_environment
 
 
 @dataclass
 class WebPageEventHandler(ApiGatewayEventHandler):
     template_name: Optional[str] = None
-    content_type: str = "text/html"
+    response_headers: Dict[str, str] = field(
+        default_factory=lambda: {"Content-Type": "text/html"}
+    )
 
     @property
     def template(self):
@@ -42,20 +45,21 @@ def handle(self, event: ExternalItemType, context) -> ExternalItemType:
         if isinstance(result, Awaitable):
             loop = asyncio.get_event_loop()
             result = loop.run_until_complete(result)
+
         if isinstance(result, Redirect):
             return {
                 "statusCode": result.status_code,
                 "headers": {"Location": result.url},
             }
 
-        dumped = self.result_marshaller.dump(result)
+        if not isinstance(result, WebPageResponse):
+            result = WebPageResponse(result, headers=self.response_headers)
+
+        dumped = self.result_marshaller.dump(result.model)
         body = self.template.render(model=dumped)
-        headers = {}
-        if self.content_type:
-            headers["Content-Type"] = self.content_type
         response = {
-            "statusCode": 200,
-            "headers": headers,
+            "statusCode": result.status_code,
+            "headers": result.headers,
             "body": body,
         }
         self.apply_caching(event, response)
@@ -97,6 +101,7 @@ def create(self, action: Action) -> Optional[WebPageEventHandler]:
             if auth_kwarg_name or action.access_control != ALLOW_ALL
             else None
         )
+        content_type = mimetypes.guess_type(action.name or "")[0] or "text/html"
         return WebPageEventHandler(
             action=action,
             param_marshaller=param_marshaller,
@@ -107,5 +112,5 @@ def create(self, action: Action) -> Optional[WebPageEventHandler]:
             authorizer=authorizer,
             priority=self.priority,
             template_name=trigger.template_name or f"{action.name}.j2",
-            content_type=mimetypes.guess_type(action.name or "")[0] or "text/html",
+            response_headers={"Content-Type": content_type},
         )

servey/servey_web_page/web_page_response.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass, field
+from typing import Generic, TypeVar, Dict
+
+T = TypeVar("T")
+
+
+@dataclass
+class WebPageResponse(Generic[T]):
+    model: T
+    status_code: int = 200
+    headers: Dict[str, str] = field(default_factory=dict)