From f22e8d3588e5b3d0b5b60ada872ab4353b2b4995 Mon Sep 17 00:00:00 2001 From: Brent Yi Date: Thu, 20 Jul 2023 14:30:51 -0700 Subject: [PATCH] Fix container removal, context thread safety --- docs/source/gui_handles.md | 15 ++++ viser/_gui_api.py | 100 +++++++++++---------- viser/client/src/ControlPanel/GuiState.tsx | 2 + viser/client/src/WebsocketInterface.tsx | 6 ++ 4 files changed, 74 insertions(+), 49 deletions(-) diff --git a/docs/source/gui_handles.md b/docs/source/gui_handles.md index 188193221..623ea9619 100644 --- a/docs/source/gui_handles.md +++ b/docs/source/gui_handles.md @@ -30,4 +30,19 @@ clients. When a GUI element is added to a client (for example, via :undoc-members: :inherited-members: +.. autoapiclass:: viser.GuiFolderHandle + :members: + :undoc-members: + :inherited-members: + +.. autoapiclass:: viser.GuiTabGroupHandle + :members: + :undoc-members: + :inherited-members: + +.. autoapiclass:: viser.GuiTabHandle + :members: + :undoc-members: + :inherited-members: + diff --git a/viser/_gui_api.py b/viser/_gui_api.py index 4e7b23851..768b1ae3d 100644 --- a/viser/_gui_api.py +++ b/viser/_gui_api.py @@ -6,12 +6,14 @@ import abc import dataclasses +import threading import time import uuid import warnings from typing import ( TYPE_CHECKING, Any, + Dict, List, Optional, Sequence, @@ -78,34 +80,35 @@ def _compute_precision_digits(x: float) -> int: return digits -# Global override for placing GUI elements in containers. -container_id_override: Optional[str] = None - -ROOT_CONTAINER_ID = "root" - - class GuiApi(abc.ABC): + _target_container_from_thread_id: Dict[int, str] = {} + """ID of container to put GUI elements into.""" + def __init__(self) -> None: super().__init__() + def _get_container_id(self) -> str: + """Get container ID associated with the current thread.""" + return self._target_container_from_thread_id.get(threading.get_ident(), "root") + + def _set_container_id(self, container_id: str) -> None: + """Set container ID associated with the current thread.""" + self._target_container_from_thread_id[threading.get_ident()] = container_id + @abc.abstractmethod def _get_api(self) -> MessageApi: """Message API to use.""" ... - def _get_container_id(self) -> str: - """ID of container to put GUI elements into.""" - if container_id_override is not None: - return container_id_override - return ROOT_CONTAINER_ID - - def gui_folder(self, label: str) -> GuiFolderHandle: - """Deprecated.""" - warnings.warn( - "gui_folder() is deprecated. Use add_gui_folder() instead!", - stacklevel=2, - ) - return self.add_gui_folder(label) + if not TYPE_CHECKING: + + def gui_folder(self, label: str) -> GuiFolderHandle: + """Deprecated.""" + warnings.warn( + "gui_folder() is deprecated. Use add_gui_folder() instead!", + stacklevel=2, + ) + return self.add_gui_folder(label) def add_gui_folder(self, label: str) -> GuiFolderHandle: """Add a folder, and return a handle that can be used to populate it.""" @@ -119,8 +122,8 @@ def add_gui_folder(self, label: str) -> GuiFolderHandle: ) ) return GuiFolderHandle( - _api=self._get_api(), - _folder_id=folder_container_id, + _gui_api=self, + _container_id=folder_container_id, ) def add_gui_tab_group(self) -> GuiTabGroupHandle: @@ -131,7 +134,7 @@ def add_gui_tab_group(self) -> GuiTabGroupHandle: _labels=[], _icons_base64=[], _tab_container_ids=[], - _api=self._get_api(), + _gui_api=self, _container_id=self._get_container_id(), ) @@ -625,7 +628,7 @@ class GuiTabGroupHandle: _labels: List[str] _icons_base64: List[Optional[str]] _tab_container_ids: List[str] - _api: MessageApi + _gui_api: GuiApi _container_id: str def add_tab(self, label: str, icon: Optional[Icon] = None) -> GuiTabHandle: @@ -644,15 +647,17 @@ def add_tab(self, label: str, icon: Optional[Icon] = None) -> GuiTabHandle: def remove(self) -> None: """Remove this tab group and all contained GUI elements.""" - self._api._queue(_messages.GuiRemoveMessage(self._tab_group_id)) - for tab_container_id in self._tab_container_ids: - self._api._queue( - _messages.GuiRemoveContainerChildrenMessage(tab_container_id) - ) + self._gui_api._get_api()._queue(_messages.GuiRemoveMessage(self._tab_group_id)) + # Containers will be removed automatically by the client. + # + # for tab_container_id in self._tab_container_ids: + # self._gui_api._get_api()._queue( + # _messages.GuiRemoveContainerChildrenMessage(tab_container_id) + # ) def _sync_with_client(self) -> None: """Send a message that syncs tab state with the client.""" - self._api._queue( + self._gui_api._get_api()._queue( _messages.GuiAddTabGroupMessage( order=time.time(), id=self._tab_group_id, @@ -668,30 +673,24 @@ def _sync_with_client(self) -> None: class GuiFolderHandle: """Use as a context to place GUI elements into a folder.""" - _api: MessageApi - _folder_id: str + _gui_api: GuiApi + _container_id: str _container_id_restore: Optional[str] = None def __enter__(self) -> None: - # For folders, we use the same ID for the container ID / folder ID. - global container_id_override - self._container_id_restore = container_id_override - container_id_override = self._folder_id + self._container_id_restore = self._gui_api._get_container_id() + self._gui_api._set_container_id(self._container_id) def __exit__(self, *args) -> None: del args - global container_id_override - container_id_override = self._container_id_restore + assert self._container_id_restore is not None + self._gui_api._set_container_id(self._container_id_restore) + self._container_id_restore = None def remove(self) -> None: """Permanently remove this folder and all contained GUI elements from the visualizer.""" - self._api._queue(_messages.GuiRemoveMessage(self._folder_id)) - # Containers corresponding to folders are removed automatically. - # - # self._parent._api._queue( - # _messages.GuiRemoveContainerChildrenMessage(self._folder_id) - # ) + self._gui_api._get_api()._queue(_messages.GuiRemoveMessage(self._container_id)) @dataclasses.dataclass @@ -703,14 +702,14 @@ class GuiTabHandle: _container_id_restore: Optional[str] = None def __enter__(self) -> None: - global container_id_override - self._container_id_restore = container_id_override - container_id_override = self._container_id + self._container_id_restore = self._parent._gui_api._get_container_id() + self._parent._gui_api._set_container_id(self._container_id) def __exit__(self, *args) -> None: del args - global container_id_override - container_id_override = self._container_id_restore + assert self._container_id_restore is not None + self._parent._gui_api._set_container_id(self._container_id_restore) + self._container_id_restore = None def remove(self) -> None: """Permanently remove this tab and all contained GUI elements from the @@ -719,7 +718,10 @@ def remove(self) -> None: container_index = self._parent._tab_container_ids.index(self._container_id) assert container_index != -1, "Tab already removed!" - print(container_index) + # Container needs to be manually removed. + self._parent._gui_api._get_api()._queue( + _messages.GuiRemoveContainerChildrenMessage(self._container_id) + ) self._parent._labels.pop(container_index) self._parent._icons_base64.pop(container_index) diff --git a/viser/client/src/ControlPanel/GuiState.tsx b/viser/client/src/ControlPanel/GuiState.tsx index 3b96a685f..ecdc731d0 100644 --- a/viser/client/src/ControlPanel/GuiState.tsx +++ b/viser/client/src/ControlPanel/GuiState.tsx @@ -108,6 +108,8 @@ export function useGuiState(initialServer: string) { const guiConfig = state.guiConfigFromId[id]; if (guiConfig.type === "GuiAddFolderMessage") state.removeGuiContainer(guiConfig.id); + if (guiConfig.type === "GuiAddTabGroupMessage") + guiConfig.tab_container_ids.forEach(state.removeGuiContainer); delete state.guiIdSetFromContainerId[guiConfig.container_id]![ guiConfig.id diff --git a/viser/client/src/WebsocketInterface.tsx b/viser/client/src/WebsocketInterface.tsx index b3c96c763..aa0a420d1 100644 --- a/viser/client/src/WebsocketInterface.tsx +++ b/viser/client/src/WebsocketInterface.tsx @@ -42,6 +42,7 @@ function useMessageHandler() { const setTheme = viewer.useGui((state) => state.setTheme); const addGui = viewer.useGui((state) => state.addGui); const removeGui = viewer.useGui((state) => state.removeGui); + const removeGuiContainer = viewer.useGui((state) => state.removeGuiContainer); const setGuiValue = viewer.useGui((state) => state.setGuiValue); const setGuiVisible = viewer.useGui((state) => state.setGuiVisible); const setGuiDisabled = viewer.useGui((state) => state.setGuiDisabled); @@ -518,6 +519,11 @@ function useMessageHandler() { removeGui(message.id); return; } + // Remove a GUI container. + case "GuiRemoveContainerChildrenMessage": { + removeGuiContainer(message.container_id); + return; + } default: { console.log("Received message did not match any known types:", message); return;