Explorar o código

split user storage into separate files; fix mypy

Falko Schindler hai 1 ano
pai
achega
bae918679d
Modificáronse 2 ficheiros con 15 adicións e 13 borrados
  1. 12 10
      nicegui/storage.py
  2. 3 3
      tests/test_storage.py

+ 12 - 10
nicegui/storage.py

@@ -5,7 +5,7 @@ import threading
 import uuid
 import uuid
 from collections.abc import MutableMapping
 from collections.abc import MutableMapping
 from pathlib import Path
 from pathlib import Path
-from typing import Any, Dict, Iterator
+from typing import Any, Dict, Iterator, Optional, Union
 
 
 import aiofiles
 import aiofiles
 from fastapi import Request
 from fastapi import Request
@@ -14,7 +14,7 @@ from starlette.responses import Response
 
 
 from . import globals
 from . import globals
 
 
-request_contextvar = contextvars.ContextVar('request_var', default=None)
+request_contextvar: contextvars.ContextVar[Optional[Request]] = contextvars.ContextVar('request_var', default=None)
 
 
 
 
 class ReadOnlyDict(MutableMapping):
 class ReadOnlyDict(MutableMapping):
@@ -97,16 +97,16 @@ class Storage:
         self.storage_dir = Path('.nicegui')
         self.storage_dir = Path('.nicegui')
         self.storage_dir.mkdir(exist_ok=True)
         self.storage_dir.mkdir(exist_ok=True)
         self._general = PersistentDict(self.storage_dir / 'storage_general.json')
         self._general = PersistentDict(self.storage_dir / 'storage_general.json')
-        self._users = PersistentDict(self.storage_dir / 'storage_users.json')
+        self._users: Dict[str, PersistentDict] = {}
 
 
     @property
     @property
-    def browser(self) -> Dict:
+    def browser(self) -> Union[ReadOnlyDict, Dict]:
         """Small storage that is saved directly within the user's browser (encrypted cookie).
         """Small storage that is saved directly within the user's browser (encrypted cookie).
 
 
         The data is shared between all browser tabs and can only be modified before the initial request has been submitted.
         The data is shared between all browser tabs and can only be modified before the initial request has been submitted.
         It is normally better to use `app.storage.user` instead to reduce payload, gain improved security and have larger storage capacity.
         It is normally better to use `app.storage.user` instead to reduce payload, gain improved security and have larger storage capacity.
         """
         """
-        request: Request = request_contextvar.get()
+        request: Optional[Request] = request_contextvar.get()
         if request is None:
         if request is None:
             if globals.get_client() == globals.index_client:
             if globals.get_client() == globals.index_client:
                 raise RuntimeError('app.storage.browser can only be used with page builder functions '
                 raise RuntimeError('app.storage.browser can only be used with page builder functions '
@@ -127,16 +127,17 @@ class Storage:
         The data is stored in a file on the server.
         The data is stored in a file on the server.
         It is shared between all browser tabs by identifying the user via session cookie ID.
         It is shared between all browser tabs by identifying the user via session cookie ID.
         """
         """
-        request: Request = request_contextvar.get()
+        request: Optional[Request] = request_contextvar.get()
         if request is None:
         if request is None:
             if globals.get_client() == globals.index_client:
             if globals.get_client() == globals.index_client:
                 raise RuntimeError('app.storage.user can only be used with page builder functions '
                 raise RuntimeError('app.storage.user can only be used with page builder functions '
                                    '(https://nicegui.io/documentation/page)')
                                    '(https://nicegui.io/documentation/page)')
             else:
             else:
                 raise RuntimeError('app.storage.user needs a storage_secret passed in ui.run()')
                 raise RuntimeError('app.storage.user needs a storage_secret passed in ui.run()')
-        if request.session['id'] not in self._users:
-            self._users[request.session['id']] = {}
-        return self._users[request.session['id']]
+        id = request.session['id']
+        if id not in self._users:
+            self._users[id] = PersistentDict(self.storage_dir / f'storage_user_{id}.json')
+        return self._users[id]
 
 
     @property
     @property
     def general(self) -> Dict:
     def general(self) -> Dict:
@@ -145,7 +146,8 @@ class Storage:
 
 
     async def backup(self) -> None:
     async def backup(self) -> None:
         await self._general.backup()
         await self._general.backup()
-        await self._users.backup()
+        for user in self._users.values():
+            await user.backup()
 
 
     async def _loop(self) -> None:
     async def _loop(self) -> None:
         while True:
         while True:

+ 3 - 3
tests/test_storage.py

@@ -86,7 +86,7 @@ async def test_access_user_storage_on_interaction(screen: Screen):
     screen.click('switch')
     screen.click('switch')
     screen.wait(1)
     screen.wait(1)
     await app.storage.backup()
     await app.storage.backup()
-    assert '{"test_switch": true}' in app.storage._users.filepath.read_text()
+    assert '{"test_switch": true}' in list(app.storage._users.values())[0].filepath.read_text()
 
 
 
 
 def test_access_user_storage_from_button_click_handler(screen: Screen):
 def test_access_user_storage_from_button_click_handler(screen: Screen):
@@ -102,7 +102,7 @@ def test_access_user_storage_from_button_click_handler(screen: Screen):
     screen.open('/')
     screen.open('/')
     screen.click('test')
     screen.click('test')
     screen.wait(1)
     screen.wait(1)
-    assert '{"inner_function": "works"}' in app.storage._users.filepath.read_text()
+    assert '{"inner_function": "works"}' in list(app.storage._users.values())[0].filepath.read_text()
 
 
 
 
 async def test_access_user_storage_from_background_task(screen: Screen):
 async def test_access_user_storage_from_background_task(screen: Screen):
@@ -116,7 +116,7 @@ async def test_access_user_storage_from_background_task(screen: Screen):
 
 
     screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.ui_run_kwargs['storage_secret'] = 'just a test'
     screen.open('/')
     screen.open('/')
-    assert '{"subtask": "works"}' in app.storage._users.filepath.read_text()
+    assert '{"subtask": "works"}' in list(app.storage._users.values())[0].filepath.read_text()
 
 
 
 
 def test_user_and_general_storage_is_persisted(screen: Screen):
 def test_user_and_general_storage_is_persisted(screen: Screen):