| from __future__ import annotations |
|
|
| import threading |
| from collections import OrderedDict |
| from copy import deepcopy |
| from typing import TYPE_CHECKING, Any |
|
|
| if TYPE_CHECKING: |
| from gradio.blocks import Blocks |
|
|
|
|
| class StateHolder: |
| def __init__(self): |
| self.capacity = 10000 |
| self.session_data = OrderedDict() |
| self.lock = threading.Lock() |
|
|
| def set_blocks(self, blocks: Blocks): |
| self.blocks = blocks |
| self.capacity = blocks.state_session_capacity |
|
|
| def __getitem__(self, session_id: str) -> SessionState: |
| if session_id not in self.session_data: |
| self.session_data[session_id] = SessionState(self.blocks) |
| self.update(session_id) |
| return self.session_data[session_id] |
|
|
| def __contains__(self, session_id: str): |
| return session_id in self.session_data |
|
|
| def update(self, session_id: str): |
| with self.lock: |
| if session_id in self.session_data: |
| self.session_data.move_to_end(session_id) |
| if len(self.session_data) > self.capacity: |
| self.session_data.popitem(last=False) |
|
|
|
|
| class SessionState: |
| def __init__(self, blocks: Blocks): |
| self.blocks = blocks |
| self._data = {} |
|
|
| def __getitem__(self, key: int) -> Any: |
| if key not in self._data: |
| block = self.blocks.blocks[key] |
| if getattr(block, "stateful", False): |
| self._data[key] = deepcopy(getattr(block, "value", None)) |
| else: |
| self._data[key] = None |
| return self._data[key] |
|
|
| def __setitem__(self, key: int, value: Any): |
| self._data[key] = value |
|
|
| def __contains__(self, key: int): |
| return key in self._data |
|
|
