update app.py

Commit 3be3692 · patx · 2025-01-27T21:28:13-05:00

Changeset
3be3692d2b7c4f8f8ac6ef14f810106667b6d1cf
Parents
69b80be427e72fbbe45b9c2f3398570cf1abe262

View source at this commit

Comments

No comments yet.

Log in to comment

Diff

diff --git a/examples/AsyncServer/WebcamStream/app.py b/examples/AsyncServer/WebcamStream/app.py
index 5db817d..5e7fa68 100644
--- a/examples/AsyncServer/WebcamStream/app.py
+++ b/examples/AsyncServer/WebcamStream/app.py
@@ -1,173 +1,80 @@
-import uvicorn
 from typing import Dict, Set, Any
-
 from MicroPie import AsyncServer
 
 # Keep track of active users (who have "started streaming")
-active_users = set()
-
-# A simple in-memory registry of connected WebSocket clients:
-# For each username, we store sets of WebSocket connections:
+active_users: Set[str] = set()
 streamers: Dict[str, Set[Any]] = {}
 watchers: Dict[str, Set[Any]] = {}
 
-
 class MyApp(AsyncServer):
 
     async def index(self):
-        """Show the home page with form to either start or watch stream."""
         return self.render_template("index.html")
 
-    async def submit(self, username, action):
-        """
-        Handle the form submission from 'index.html'.
-        Allows user to pick 'Start Streaming' or 'Watch Stream'.
-        """
+    async def submit(self, username: str, action: str):
         if username:
             active_users.add(username)
-            if action == "Start Streaming":
-                return self.redirect(f"/stream/{username}")
-            elif action == "Watch Stream":
-                return self.redirect(f"/watch/{username}")
+            route = f"/stream/{username}" if action == "Start Streaming" else f"/watch/{username}"
+            return self.redirect(route)
         return self.redirect("/")
 
-    async def stream(self, username):
-        """
-        Show the streaming page for a particular username.
-        Only valid if the username is in active_users.
-        """
-        if username not in active_users:
-            return self.redirect("/")
-        return self.render_template("stream.html", username=username)
-
-    async def watch(self, username):
-        """
-        Show the watch page for a particular username.
-        Only valid if the username is in active_users.
-        """
-        if username not in active_users:
-            return self.redirect("/")
-        return self.render_template("watch.html", username=username)
-
-    #
-    # WebSocket routes
-    #
+    async def stream(self, username: str):
+        return self.render_template("stream.html", username=username) if username in active_users else self.redirect("/")
+
+    async def watch(self, username: str):
+        return self.render_template("watch.html", username=username) if username in active_users else self.redirect("/")
+
     async def websocket_stream(self, scope, receive, send):
-        """
-        Handle WebSocket connections for a streamer, i.e. /stream/<username>.
-        Path params: [username]
-        """
-        # We expect self.path_params[0] to be the username.
-        if not self.path_params:
-            await self._reject_websocket(send)
-            return
-
-        username = self.path_params[0]
-        if username not in active_users:
-            await self._reject_websocket(send)
-            return
-
-        # Accept the WebSocket connection
+        username = self.path_params[0] if self.path_params else None
+        if not username or username not in active_users:
+            return await self._reject_websocket(send)
+
         await send({"type": "websocket.accept"})
+        streamers.setdefault(username, set()).add(send)
 
-        # Ensure we have a set for this username
-        if username not in streamers:
-            streamers[username] = set()
-        streamers[username].add(send)
-
-        try:
-            while True:
-                message = await receive()
-                if message["type"] == "websocket.receive":
-                    # We expect JSON or text frames. If binary, adapt as needed.
-                    # For a simple approach, let's assume 'text' with a JSON-like structure
-                    # or a raw image frame in 'bytes'.
-                    if "text" in message:
-                        text_data = message["text"]
-                        # Broadcast text_data to watchers of this username
-                        await self._broadcast(username, text_data)
-                    elif "bytes" in message:
-                        binary_frame = message["bytes"]
-                        # Broadcast binary frame to watchers of this username
-                        await self._broadcast(username, binary_frame, is_binary=True)
-                elif message["type"] == "websocket.disconnect":
-                    break
-        finally:
-            # On disconnect, remove this send channel from the streamer set
-            streamers[username].discard(send)
-            if not streamers[username]:
-                del streamers[username]
+        await self._handle_websocket(receive, send, username, streamers, is_stream=True)
 
     async def websocket_watch(self, scope, receive, send):
-        """
-        Handle WebSocket connections for watchers, i.e. /watch/<username>.
-        Path params: [username]
-        """
-        # We expect self.path_params[0] to be the username.
-        if not self.path_params:
-            await self._reject_websocket(send)
-            return
-
-        username = self.path_params[0]
-        if username not in active_users:
-            await self._reject_websocket(send)
-            return
-
-        # Accept the WebSocket connection
+        username = self.path_params[0] if self.path_params else None
+        if not username or username not in active_users:
+            return await self._reject_websocket(send)
+
         await send({"type": "websocket.accept"})
+        watchers.setdefault(username, set()).add(send)
 
-        # Ensure we have a set for watchers of this username
-        if username not in watchers:
-            watchers[username] = set()
-        watchers[username].add(send)
+        await self._handle_websocket(receive, send, username, watchers)
 
+    async def _handle_websocket(self, receive, send, username, registry, is_stream=False):
         try:
             while True:
                 message = await receive()
-                # Typically watchers won't send frames, but you could handle
-                # text commands from watchers if needed.
-                if message["type"] == "websocket.disconnect":
+                if message["type"] == "websocket.receive":
+                    if "text" in message or "bytes" in message:
+                        await self._broadcast(username, message.get("text") or message.get("bytes"), is_binary="bytes" in message)
+                elif message["type"] == "websocket.disconnect":
                     break
         finally:
-            # On disconnect, remove this send channel from the watchers set
-            watchers[username].discard(send)
-            if not watchers[username]:
-                del watchers[username]
-
-    #
-    # Helper methods
-    #
+            registry[username].discard(send)
+            if not registry[username]:
+                del registry[username]
+
     async def _reject_websocket(self, send):
-        """Reject a WebSocket connection (invalid user or path)."""
         await send({"type": "websocket.accept"})
         await send({"type": "websocket.close", "code": 4000})
 
     async def _broadcast(self, username: str, data, is_binary: bool = False):
-        """
-        Broadcast a message to all watchers of the specified username.
-        data can be str or bytes. is_binary toggles send mode.
-        """
-        if username not in watchers:
-            return
-        for ws_send in list(watchers[username]):
+        for ws_send in list(watchers.get(username, [])):
             try:
-                if is_binary:
-                    await ws_send({"type": "websocket.send", "bytes": data})
-                else:
-                    await ws_send({"type": "websocket.send", "text": data})
+                await ws_send({"type": "websocket.send", "bytes" if is_binary else "text": data})
             except:
-                # If we fail to send, assume the connection is dead
                 watchers[username].discard(ws_send)
                 if not watchers[username]:
                     del watchers[username]
 
-
-# Create our application instance
+# Create the ASGI app
 app = MyApp()
 
-#
-# If you want to run directly from python (instead of `uvicorn myapp:app`), do:
-#
 if __name__ == "__main__":
-    uvicorn.run("myapp:app", host="0.0.0.0", port=5000, reload=True)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=5000, reload=True)