Browse Source

introduce single-use routes for ui.download

Falko Schindler 1 year ago
parent
commit
da356b8c18
3 changed files with 20 additions and 4 deletions
  1. 16 2
      nicegui/app.py
  2. 1 1
      nicegui/functions/download.py
  3. 3 1
      tests/test_download.py

+ 16 - 2
nicegui/app.py

@@ -82,7 +82,11 @@ class App(FastAPI):
             raise ValueError('''Path cannot be "/", because it would hide NiceGUI's internal "/_nicegui" route.''')
         globals.app.mount(url_path, StaticFiles(directory=str(local_directory)))
 
-    def add_static_file(self, *, local_file: Union[str, Path], url_path: Optional[str] = None) -> str:
+    def add_static_file(self, *,
+                        local_file: Union[str, Path],
+                        url_path: Optional[str] = None,
+                        single_use: bool = False,
+                        ) -> str:
         """Add a single static file.
 
         Allows a local file to be accessed online with enabled caching.
@@ -93,6 +97,7 @@ class App(FastAPI):
 
         :param local_file: local file to serve as static content
         :param url_path: string that starts with a slash "/" and identifies the path at which the file should be served (default: None -> auto-generated URL path)
+        :param single_use: whether to remove the route after the file has been downloaded once (default: False)
         :return: URL path which can be used to access the file
         """
         file = Path(local_file).resolve()
@@ -103,6 +108,8 @@ class App(FastAPI):
 
         @self.get(url_path)
         def read_item() -> FileResponse:
+            if single_use:
+                self.remove_route(url_path)
             return FileResponse(file, headers={'Cache-Control': 'public, max-age=3600'})
 
         return url_path
@@ -128,7 +135,11 @@ class App(FastAPI):
                 return {'detail': 'Not Found'}, 404
             return helpers.get_streaming_response(filepath, request)
 
-    def add_media_file(self, *, local_file: Union[str, Path], url_path: Optional[str] = None) -> str:
+    def add_media_file(self, *,
+                       local_file: Union[str, Path],
+                       url_path: Optional[str] = None,
+                       single_use: bool = False,
+                       ) -> str:
         """Add a single media file.
 
         Allows a local file to be streamed.
@@ -139,6 +150,7 @@ class App(FastAPI):
 
         :param local_file: local file to serve as media content
         :param url_path: string that starts with a slash "/" and identifies the path at which the file should be served (default: None -> auto-generated URL path)
+        :param single_use: whether to remove the route after the media file has been downloaded once (default: False)
         :return: URL path which can be used to access the file
         """
         file = Path(local_file).resolve()
@@ -149,6 +161,8 @@ class App(FastAPI):
 
         @self.get(url_path)
         def read_item(request: Request) -> StreamingResponse:
+            if single_use:
+                self.remove_route(url_path)
             return helpers.get_streaming_response(file, request)
 
         return url_path

+ 1 - 1
nicegui/functions/download.py

@@ -13,5 +13,5 @@ def download(src: Union[str, Path], filename: Optional[str] = None) -> None:
     :param filename: name of the file to download (default: name of the file on the server)
     """
     if helpers.is_file(src):
-        src = globals.app.add_static_file(local_file=src)
+        src = globals.app.add_static_file(local_file=src, single_use=True)
     globals.get_client().download(src, filename)

+ 3 - 1
tests/test_download.py

@@ -12,7 +12,7 @@ from .screen import Screen
 
 @pytest.fixture
 def test_route() -> Generator[str, None, None]:
-    TEST_ROUTE = '/static/test.py'
+    TEST_ROUTE = '/static/test.txt'
     yield TEST_ROUTE
     app.remove_route(TEST_ROUTE)
 
@@ -35,6 +35,8 @@ def test_downloading_local_file_as_src(screen: Screen):
     ui.button('download', on_click=lambda: ui.download(IMAGE_FILE))
 
     screen.open('/')
+    route_count_before_download = len(app.routes)
     screen.click('download')
     screen.wait(0.5)
     assert (DOWNLOAD_DIR / 'slide1.jpg').exists()
+    assert len(app.routes) == route_count_before_download