Merge pull request #2250 from zauberzeug/on_air_media_files

Allow http range requests with On Air (eg. serving media files)

nicegui/air.py

@@ -2,7 +2,9 @@ import asyncio
 import gzip
 import json
 import re
-from typing import Any, Dict, Optional
+from dataclasses import dataclass
+from typing import Any, AsyncIterator, Dict, Optional
+from uuid import uuid4
 
 import httpx
 import socketio
@@ -10,18 +12,28 @@ import socketio.exceptions
 
 from . import background_tasks, core
 from .client import Client
+from .dataclasses import KWONLY_SLOTS
 from .logging import log
 
 RELAY_HOST = 'https://on-air.nicegui.io/'
 
 
+@dataclass(**KWONLY_SLOTS)
+class Stream:
+    data: AsyncIterator[bytes]
+    response: httpx.Response
+
+
 class Air:
 
     def __init__(self, token: str) -> None:
         self.token = token
         self.relay = socketio.AsyncClient()
         self.client = httpx.AsyncClient(app=core.app)
+        self.streaming_client = httpx.AsyncClient()
         self.connecting = False
+        self.streams: Dict[str, Stream] = {}
+        self.remote_url: Optional[str] = None
 
         @self.relay.on('http')
         async def _handle_http(data: Dict[str, Any]) -> Dict[str, Any]:
@@ -53,9 +65,46 @@ class Air:
                 'content': compressed,
             }
 
+        @self.relay.on('range-request')
+        async def _handle_range_request(data: Dict[str, Any]) -> Dict[str, Any]:
+            headers: Dict[str, Any] = data['headers']
+            url = list(u for u in core.app.urls if self.remote_url != u)[0] + data['path']
+            data['params']['nicegui_chunk_size'] = 1024
+            request = self.client.build_request(
+                data['method'],
+                url,
+                params=data['params'],
+                headers=headers,
+            )
+            response = await self.streaming_client.send(request, stream=True)
+            stream_id = str(uuid4())
+            self.streams[stream_id] = Stream(data=response.aiter_bytes(), response=response)
+            return {
+                'status_code': response.status_code,
+                'headers': response.headers.multi_items(),
+                'stream_id': stream_id,
+            }
+
+        @self.relay.on('read-stream')
+        async def _handle_read_stream(stream_id: str) -> Optional[bytes]:
+            try:
+                return await self.streams[stream_id].data.__anext__()
+            except StopAsyncIteration:
+                await _handle_close_stream(stream_id)
+                return None
+            except Exception:
+                await _handle_close_stream(stream_id)
+                raise
+
+        @self.relay.on('close-stream')
+        async def _handle_close_stream(stream_id: str) -> None:
+            await self.streams[stream_id].response.aclose()
+            del self.streams[stream_id]
+
         @self.relay.on('ready')
         def _handle_ready(data: Dict[str, Any]) -> None:
             core.app.urls.add(data['device_url'])
+            self.remote_url = data['device_url']
             if core.app.config.show_welcome_message:
                 print(f'NiceGUI is on air at {data["device_url"]}', flush=True)
 
@@ -139,6 +188,9 @@ class Air:
 
     async def disconnect(self) -> None:
         """Disconnect from the NiceGUI On Air server."""
+        for stream in self.streams.values():
+            await stream.response.aclose()
+        self.streams.clear()
         await self.relay.disconnect()
 
     async def emit(self, message_type: str, data: Dict[str, Any], room: str) -> None:

nicegui/app/app.py

@@ -3,8 +3,8 @@ from enum import Enum
 from pathlib import Path
 from typing import Any, Awaitable, Callable, List, Optional, Union
 
-from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import FileResponse, StreamingResponse
+from fastapi import FastAPI, HTTPException, Request, Response
+from fastapi.responses import FileResponse
 from fastapi.staticfiles import StaticFiles
 
 from .. import background_tasks, helpers
@@ -15,7 +15,7 @@ from ..observables import ObservableSet
 from ..server import Server
 from ..storage import Storage
 from .app_config import AppConfig
-from .streaming_response import get_streaming_response
+from .range_response import get_range_response
 
 
 class State(Enum):
@@ -196,11 +196,11 @@ class App(FastAPI):
         :param local_directory: local folder with files to serve as media content
         """
         @self.get(url_path + '/{filename:path}')
-        def read_item(request: Request, filename: str) -> StreamingResponse:
+        def read_item(request: Request, filename: str, nicegui_cunk_size: int = 8192) -> Response:
             filepath = Path(local_directory) / filename
             if not filepath.is_file():
                 raise HTTPException(status_code=404, detail='Not Found')
-            return get_streaming_response(filepath, request)
+            return get_range_response(filepath, request, chunk_size=nicegui_cunk_size)
 
     def add_media_file(self, *,
                        local_file: Union[str, Path],
@@ -226,10 +226,10 @@ class App(FastAPI):
         path = f'/_nicegui/auto/media/{helpers.hash_file_path(file)}/{file.name}' if url_path is None else url_path
 
         @self.get(path)
-        def read_item(request: Request) -> StreamingResponse:
+        def read_item(request: Request, nicegui_cunk_size: int = 8192) -> Response:
             if single_use:
                 self.remove_route(path)
-            return get_streaming_response(file, request)
+            return get_range_response(file, request, chunk_size=nicegui_cunk_size)
 
         return path
 

nicegui/app/range_response.py

@@ -0,0 +1,58 @@
+import hashlib
+import mimetypes
+from datetime import datetime
+from pathlib import Path
+from typing import Generator
+
+from fastapi import Request
+from fastapi.responses import Response, StreamingResponse
+
+mimetypes.init()
+
+
+def get_range_response(file: Path, request: Request, chunk_size: int) -> Response:
+    """Get a Response for the given file, supporting range-requests, E-Tag and Last-Modified."""
+    file_size = file.stat().st_size
+    last_modified_time = datetime.utcfromtimestamp(file.stat().st_mtime)
+    start = 0
+    end = file_size - 1
+    status_code = 200
+    e_tag = hashlib.md5((str(last_modified_time) + str(file_size)).encode()).hexdigest()
+    if_match_header = request.headers.get('If-None-Match')
+    if if_match_header and if_match_header == e_tag:
+        return Response(status_code=304)  # Not Modified
+    headers = {
+        'E-Tag': e_tag,
+        'Last-Modified': last_modified_time.strftime(r'%a, %d %b %Y %H:%M:%S GMT'),
+    }
+    range_header = request.headers.get('range')
+    media_type = mimetypes.guess_type(str(file))[0] or 'application/octet-stream'
+    if range_header is not None:
+        byte1, byte2 = range_header.split('=')[1].split('-')
+        start = int(byte1)
+        if byte2:
+            end = int(byte2)
+        status_code = 206  # Partial Content
+    content_length = end - start + 1
+    headers.update({
+        'Content-Length': str(content_length),
+        'Content-Range': f'bytes {start}-{end}/{file_size}',
+        'Accept-Ranges': 'bytes',
+    })
+
+    def content_reader(file: Path, start: int, end: int) -> Generator[bytes, None, None]:
+        with open(file, 'rb') as data:
+            data.seek(start)
+            remaining_bytes = end - start + 1
+            while remaining_bytes > 0:
+                chunk = data.read(min(chunk_size, remaining_bytes))
+                if not chunk:
+                    break
+                yield chunk
+                remaining_bytes -= len(chunk)
+    return StreamingResponse(
+        content_reader(file, start, end),
+        media_type=media_type,
+        headers=headers,
+        status_code=status_code,
+    )

nicegui/app/streaming_response.py

@@ -1,45 +0,0 @@
-import mimetypes
-from pathlib import Path
-from typing import Generator
-
-from fastapi import Request
-from fastapi.responses import StreamingResponse
-
-mimetypes.init()
-
-
-def get_streaming_response(file: Path, request: Request) -> StreamingResponse:
-    """Get a StreamingResponse for the given file and request."""
-    file_size = file.stat().st_size
-    start = 0
-    end = file_size - 1
-    range_header = request.headers.get('Range')
-    if range_header:
-        byte1, byte2 = range_header.split('=')[1].split('-')
-        start = int(byte1)
-        if byte2:
-            end = int(byte2)
-    content_length = end - start + 1
-    headers = {
-        'Content-Range': f'bytes {start}-{end}/{file_size}',
-        'Content-Length': str(content_length),
-        'Accept-Ranges': 'bytes',
-    }
-
-    def content_reader(file: Path, start: int, end: int, chunk_size: int = 8192) -> Generator[bytes, None, None]:
-        with open(file, 'rb') as data:
-            data.seek(start)
-            remaining_bytes = end - start + 1
-            while remaining_bytes > 0:
-                chunk = data.read(min(chunk_size, remaining_bytes))
-                if not chunk:
-                    break
-                yield chunk
-                remaining_bytes -= len(chunk)
-
-    return StreamingResponse(
-        content_reader(file, start, end),
-        media_type=mimetypes.guess_type(str(file))[0] or 'application/octet-stream',
-        headers=headers,
-        status_code=206,
-    )

nicegui/nicegui.py

@@ -121,9 +121,9 @@ async def _shutdown() -> None:
     """Handle the shutdown event."""
     if app.native.main_window:
         app.native.main_window.signal_server_shutdown()
+    air.disconnect()
     app.stop()
     run.tear_down()
-    air.disconnect()
 
 
 @app.exception_handler(404)

nicegui/templates/index.html

@@ -214,7 +214,12 @@
 
       function download(src, filename) {
         const anchor = document.createElement("a");
-        anchor.href = typeof src === "string" ? src : URL.createObjectURL(new Blob([src]));
+        if (typeof src === "string") {
+          anchor.href = src.startsWith("/") ? "{{ prefix | safe }}" + src : src;
+        }
+        else {
+          anchor.href = URL.createObjectURL(new Blob([src]))
+        }
         anchor.target = "_blank";
         anchor.download = filename || "";
         document.body.appendChild(anchor);