Add support for response generators

Author: jtc42

webthing/server.py

@@ -4,6 +4,9 @@
 import json
 import socket
 import sys
+import typing
+import types
+import asyncio
 import tornado.concurrent
 import tornado.gen
 import tornado.httpserver
@@ -56,6 +59,39 @@ def options(self, *args, **kwargs):
         """Handle an OPTIONS request."""
         self.set_status(204)
 
+    async def represent_response(
+        self, data, content_type: str = "application/json", headers: dict = None
+    ):
+        headers = headers or {}
+        for k, v in headers:
+            self.set_header(k, v)
+
+        self.set_header("Content-Type", content_type)
+
+        if isinstance(data, (typing.AsyncGenerator, types.GeneratorType)):
+            self.set_header(
+                "Cache-Control",
+                "no-store, no-cache, must-revalidate, pre-check=0, post-check=0, max-age=0",
+            )
+            self.set_header("Pragma", "no-cache")
+
+            if isinstance(data, typing.AsyncGenerator):
+                async for frame in data:
+                    # Write data to memory
+                    self.write(frame)
+                    # Write data to network
+                    await self.flush()
+            elif isinstance(data, types.GeneratorType):
+                for frame in data:
+                    # Write data to memory
+                    self.write(frame)
+                    # Write data to network
+                    await self.flush()
+        else:
+            if content_type == "application/json":
+                data = json.dumps(data)
+            self.write(data)
+
 
 class ThingHandler(tornado.websocket.WebSocketHandler, Subscriber):
     """Handle a request to /."""
@@ -299,7 +335,7 @@ def update_event(self, event):
 class PropertiesHandler(BaseHandler):
     """Handle a request to /properties."""
 
-    def get(self):
+    async def get(self):
         """
         Handle a GET request.
         """
@@ -308,13 +344,13 @@ def get(self):
             return
 
         self.set_header("Content-Type", "application/json")
-        self.write(json.dumps(self.thing.get_properties()))
+        self.write(json.dumps(await self.thing.get_properties()))
 
 
 class PropertyHandler(BaseHandler):
     """Handle a request to /properties/<property>."""
 
-    def get(self, property_name=None):
+    async def get(self, property_name=None):
         """
         Handle a GET request.
 
@@ -324,13 +360,14 @@ def get(self, property_name=None):
             self.set_status(404)
             return
 
-        if self.thing.has_property(property_name):
-            self.set_header("Content-Type", "application/json")
-            self.write(json.dumps(self.thing.get_property(property_name)))
+        prop = self.thing.find_property(property_name)
+        if prop:
+            data = await self.thing.get_property(property_name)
+            await self.represent_response(data, prop.get_content_type())
         else:
             self.set_status(404)
 
-    def put(self, property_name=None):
+    async def put(self, property_name=None):
         """
         Handle a PUT request.
 
@@ -346,20 +383,15 @@ def put(self, property_name=None):
             self.set_status(400)
             return
 
-        if self.thing.has_property(property_name):
+        prop = self.thing.find_property(property_name)
+        if prop:
             try:
-                self.thing.set_property(property_name, args)
+                await self.thing.set_property(property_name, args)
             except PropertyError:
                 self.set_status(400)
                 return
-
-            self.set_header("Content-Type", "application/json")
-            self.write(
-                json.dumps(
-                    {
-                        property_name: self.thing.get_property(property_name),
-                    }
-                )
+            await self.represent_response(
+                await self.thing.get_property(property_name), prop.get_content_type()
             )
         else:
             self.set_status(404)
@@ -368,18 +400,17 @@ def put(self, property_name=None):
 class ActionsHandler(BaseHandler):
     """Handle a request to /actions."""
 
-    def get(self):
+    async def get(self):
         """
         Handle a GET request.
         """
         if self.thing is None:
             self.set_status(404)
             return
 
-        self.set_header("Content-Type", "application/json")
-        self.write(json.dumps(self.thing.get_action_descriptions()))
+        await self.represent_response(self.thing.get_action_descriptions())
 
-    def post(self):
+    async def post(self):
         """
         Handle a POST request.
         """
@@ -412,15 +443,15 @@ def post(self):
             tornado.ioloop.IOLoop.current().spawn_callback(action.start)
 
             self.set_status(201)
-            self.write(json.dumps(response))
+            await self.represent_response(response)
         else:
             self.set_status(400)
 
 
 class ActionHandler(BaseHandler):
     """Handle a request to /actions/<action_name>."""
 
-    def get(self, action_name=None):
+    async def get(self, action_name=None):
         """
         Handle a GET request.
 
@@ -430,10 +461,11 @@ def get(self, action_name=None):
             self.set_status(404)
             return
 
-        self.set_header("Content-Type", "application/json")
-        self.write(json.dumps(self.thing.get_action_descriptions(action_name=action_name)))
+        await self.represent_response(
+            self.thing.get_action_descriptions(action_name=action_name)
+        )
 
-    def post(self, action_name=None):
+    async def post(self, action_name=None):
         """
         Handle a POST request.
         """
@@ -455,15 +487,15 @@ def post(self, action_name=None):
             tornado.ioloop.IOLoop.current().spawn_callback(action.start)
 
             self.set_status(201)
-            self.write(json.dumps(response))
+            await self.represent_response(response)
         else:
             self.set_status(400)
 
 
 class ActionIDHandler(BaseHandler):
     """Handle a request to /actions/<action_name>/<action_id>."""
 
-    def get(self, action_name=None, action_id=None):
+    async def get(self, action_name=None, action_id=None):
         """
         Handle a GET request.
 
@@ -479,10 +511,9 @@ def get(self, action_name=None, action_id=None):
             self.set_status(404)
             return
 
-        self.set_header("Content-Type", "application/json")
-        self.write(json.dumps(action.as_action_description()))
+        await self.represent_response(action.as_action_description())
 
-    def put(self,  action_name=None, action_id=None):
+    def put(self, action_name=None, action_id=None):
         """
         Handle a PUT request.
 
@@ -517,22 +548,21 @@ def delete(self, action_name=None, action_id=None):
 class EventsHandler(BaseHandler):
     """Handle a request to /events."""
 
-    def get(self):
+    async def get(self):
         """
         Handle a GET request.
         """
         if self.thing is None:
             self.set_status(404)
             return
 
-        self.set_header("Content-Type", "application/json")
-        self.write(json.dumps(self.thing.get_event_descriptions()))
+        await self.represent_response(self.thing.get_event_descriptions())
 
 
 class EventHandler(BaseHandler):
     """Handle a request to /events/<event_name>."""
 
-    def get(self, event_name=None):
+    async def get(self, event_name=None):
         """
         Handle a GET request.
 
@@ -542,8 +572,9 @@ def get(self, event_name=None):
             self.set_status(404)
             return
 
-        self.set_header("Content-Type", "application/json")
-        self.write(json.dumps(self.thing.get_event_descriptions(event_name=event_name)))
+        await self.represent_response(
+            self.thing.get_event_descriptions(event_name=event_name)
+        )
 
 
 class WebThingServer: