Add progressbar component (#238)

* Add progressbar component * Minor nits: add_progressbar => add_progress_bar, match Mantine naming, CSS tweak * Formatting * Comment nits * ruff --------- Co-authored-by: Brent Yi <[email protected]>
nerfstudio-project · Jul 16, 2024 · fbcb430 · fbcb430
1 parent 03ce7f7
commit fbcb430
Show file tree

Hide file tree

Showing 8 changed files with 235 additions and 76 deletions.
diff --git a/.gitignore b/.gitignore
@@ -14,5 +14,6 @@ htmlcov
 .DS_Store
 .envrc
 .vite
+build
 src/viser/client/build
 src/viser/client/.nodeenv
diff --git a/examples/02_gui.py b/examples/02_gui.py
@@ -18,7 +18,6 @@ def main() -> None:
             initial_value=0,
             disabled=True,
         )
-
         gui_slider = server.gui.add_slider(
             "Slider",
             min=0,
@@ -27,6 +26,7 @@ def main() -> None:
             initial_value=0,
             disabled=True,
         )
+        gui_progress = server.gui.add_progress_bar(25, animated=True)
 
     with server.gui.add_folder("Editable"):
         gui_vector2 = server.gui.add_vector2(
@@ -108,6 +108,8 @@ def _(_) -> None:
             point_shape="circle",
         )
 
+        gui_progress.value = float((counter % 100))
+
         # We can use `.visible` and `.disabled` to toggle GUI elements.
         gui_text.visible = not gui_checkbox_hide.value
         gui_button.visible = not gui_checkbox_hide.value

diff --git a/src/viser/_gui_api.py b/src/viser/_gui_api.py
@@ -32,6 +32,7 @@
     GuiMarkdownHandle,
     GuiModalHandle,
     GuiPlotlyHandle,
+    GuiProgressBarHandle,
     GuiTabGroupHandle,
     GuiUploadButtonHandle,
     SupportsRemoveProtocol,
@@ -57,6 +58,22 @@
 TLiteralString = TypeVar("TLiteralString", bound=LiteralString)
 T = TypeVar("T")
 LengthTenStrTuple: TypeAlias = Tuple[str, str, str, str, str, str, str, str, str, str]
+Color: TypeAlias = Literal[
+    "dark",
+    "gray",
+    "red",
+    "pink",
+    "grape",
+    "violet",
+    "indigo",
+    "blue",
+    "cyan",
+    "green",
+    "lime",
+    "yellow",
+    "orange",
+    "teal",
+]
 
 
 def _hex_from_hls(h: float, l: float, s: float) -> str:
@@ -637,23 +654,7 @@ def add_button(
         disabled: bool = False,
         visible: bool = True,
         hint: str | None = None,
-        color: Literal[
-            "dark",
-            "gray",
-            "red",
-            "pink",
-            "grape",
-            "violet",
-            "indigo",
-            "blue",
-            "cyan",
-            "green",
-            "lime",
-            "yellow",
-            "orange",
-            "teal",
-        ]
-        | None = None,
+        color: Color | None = None,
         icon: IconName | None = None,
         order: float | None = None,
     ) -> GuiButtonHandle:
@@ -701,23 +702,7 @@ def add_upload_button(
         disabled: bool = False,
         visible: bool = True,
         hint: str | None = None,
-        color: Literal[
-            "dark",
-            "gray",
-            "red",
-            "pink",
-            "grape",
-            "violet",
-            "indigo",
-            "blue",
-            "cyan",
-            "green",
-            "lime",
-            "yellow",
-            "orange",
-            "teal",
-        ]
-        | None = None,
+        color: Color | None = None,
         icon: IconName | None = None,
         mime_type: str = "*/*",
         order: float | None = None,
@@ -1176,6 +1161,49 @@ def add_dropdown(
             _impl_options=tuple(options),
         )
 
+    def add_progress_bar(
+        self,
+        value: float,
+        visible: bool = True,
+        animated: bool = False,
+        color: Color | None = None,
+        order: float | None = None,
+    ) -> GuiProgressBarHandle:
+        """Add a progress bar to the GUI.
+
+        Args:
+            value: Value of the progress bar. (0 - 100)
+            visible: Whether the progress bar is visible.
+            animated: Whether the progress bar is in a loading state (animated, striped).
+            color: The color of the progress bar.
+            order: Optional ordering, smallest values will be displayed first.
+
+        Returns:
+            A handle that can be used to interact with the GUI element.
+        """
+        assert value >= 0 and value <= 100
+        handle = GuiProgressBarHandle(
+            _gui_api=self,
+            _id=_make_unique_id(),
+            _visible=visible,
+            _animated=animated,
+            _parent_container_id=self._get_container_id(),
+            _order=_apply_default_order(order),
+            _value=value,
+        )
+        self._websock_interface.queue_message(
+            _messages.GuiAddProgressBarMessage(
+                order=handle._order,
+                id=handle._id,
+                value=value,
+                animated=animated,
+                color=color,
+                container_id=handle._parent_container_id,
+                visible=visible,
+            )
+        )
+        return handle
+
     def add_slider(
         self,
         label: str,

diff --git a/src/viser/_gui_handles.py b/src/viser/_gui_handles.py
@@ -559,14 +559,92 @@ def _parse_markdown(markdown: str, image_root: Path | None) -> str:
     return markdown
 
 
+@dataclasses.dataclass
+class GuiProgressBarHandle:
+    """Use to remove markdown."""
+
+    _gui_api: GuiApi
+    _id: str
+    _visible: bool
+    _animated: bool
+    _parent_container_id: str
+    _order: float
+    _value: float
+
+    @property
+    def value(self) -> float:
+        """Current content of this progress bar element, 0 - 100. Synchronized
+        automatically when assigned."""
+        return self._value
+
+    @value.setter
+    def value(self, value: float) -> None:
+        assert value >= 0 and value <= 100
+        self._value = value
+        self._gui_api._websock_interface.queue_message(
+            GuiUpdateMessage(
+                self._id,
+                {"value": value},
+            )
+        )
+
+    @property
+    def animated(self) -> bool:
+        """Show this progress bar as loading (animated, striped)."""
+        return self._animated
+
+    @animated.setter
+    def animated(self, animated: bool) -> None:
+        self._animated = animated
+        self._gui_api._websock_interface.queue_message(
+            GuiUpdateMessage(
+                self._id,
+                {"animated": animated},
+            )
+        )
+
+    @property
+    def order(self) -> float:
+        """Read-only order value, which dictates the position of the GUI element."""
+        return self._order
+
+    @property
+    def visible(self) -> bool:
+        """Temporarily show or hide this GUI element from the visualizer. Synchronized
+        automatically when assigned."""
+        return self._visible
+
+    @visible.setter
+    def visible(self, visible: bool) -> None:
+        if visible == self.visible:
+            return
+
+        self._gui_api._websock_interface.queue_message(
+            GuiUpdateMessage(self._id, {"visible": visible})
+        )
+        self._visible = visible
+
+    def __post_init__(self) -> None:
+        """We need to register ourself after construction for callbacks to work."""
+        parent = self._gui_api._container_handle_from_id[self._parent_container_id]
+        parent._children[self._id] = self
+
+    def remove(self) -> None:
+        """Permanently remove this progress bar from the visualizer."""
+        self._gui_api._websock_interface.queue_message(GuiRemoveMessage(self._id))
+
+        parent = self._gui_api._container_handle_from_id[self._parent_container_id]
+        parent._children.pop(self._id)
+
+
 @dataclasses.dataclass
 class GuiMarkdownHandle:
     """Use to remove markdown."""
 
     _gui_api: GuiApi
     _id: str
     _visible: bool
-    _parent_container_id: str  # Parent.
+    _parent_container_id: str
     _order: float
     _image_root: Path | None
     _content: str | None
@@ -628,7 +706,7 @@ class GuiPlotlyHandle:
     _gui_api: GuiApi
     _id: str
     _visible: bool
-    _parent_container_id: str  # Parent.
+    _parent_container_id: str
     _order: float
     _figure: go.Figure | None
     _aspect: float | None
@@ -696,7 +774,7 @@ def __post_init__(self) -> None:
         parent._children[self._id] = self
 
     def remove(self) -> None:
-        """Permanently remove this markdown from the visualizer."""
+        """Permanently remove this figure from the visualizer."""
         self._gui_api._websock_interface.queue_message(GuiRemoveMessage(self._id))
         parent = self._gui_api._container_handle_from_id[self._parent_container_id]
         parent._children.pop(self._id)
diff --git a/src/viser/_messages.py b/src/viser/_messages.py
@@ -25,6 +25,22 @@
 from . import infra, theme
 
 GuiSliderMark = TypedDict("GuiSliderMark", {"value": float, "label": NotRequired[str]})
+Color = Literal[
+    "dark",
+    "gray",
+    "red",
+    "pink",
+    "grape",
+    "violet",
+    "indigo",
+    "blue",
+    "cyan",
+    "green",
+    "lime",
+    "yellow",
+    "orange",
+    "teal",
+]
 
 
 class Message(infra.Message):
@@ -483,6 +499,18 @@ class GuiAddMarkdownMessage(Message):
     visible: bool
 
 
+@tag_class("GuiAddComponentMessage")
+@dataclasses.dataclass
+class GuiAddProgressBarMessage(Message):
+    order: float
+    id: str
+    value: float
+    animated: bool
+    color: Optional[Color]
+    container_id: str
+    visible: bool
+
+
 @tag_class("GuiAddComponentMessage")
 @dataclasses.dataclass
 class GuiAddPlotlyMessage(Message):
@@ -538,48 +566,14 @@ class GuiAddButtonMessage(_GuiAddInputBase):
     # All GUI elements currently need an `value` field.
     # This makes our job on the frontend easier.
     value: bool
-    color: Optional[
-        Literal[
-            "dark",
-            "gray",
-            "red",
-            "pink",
-            "grape",
-            "violet",
-            "indigo",
-            "blue",
-            "cyan",
-            "green",
-            "lime",
-            "yellow",
-            "orange",
-            "teal",
-        ]
-    ]
+    color: Optional[Color]
     icon_html: Optional[str]
 
 
 @tag_class("GuiAddComponentMessage")
 @dataclasses.dataclass
 class GuiAddUploadButtonMessage(_GuiAddInputBase):
-    color: Optional[
-        Literal[
-            "dark",
-            "gray",
-            "red",
-            "pink",
-            "grape",
-            "violet",
-            "indigo",
-            "blue",
-            "cyan",
-            "green",
-            "lime",
-            "yellow",
-            "orange",
-            "teal",
-        ]
-    ]
+    color: Optional[Color]
     icon_html: Optional[str]
     mime_type: str
 

diff --git a/src/viser/client/src/ControlPanel/Generated.tsx b/src/viser/client/src/ControlPanel/Generated.tsx
@@ -21,6 +21,7 @@ import TabGroupComponent from "../components/TabGroup";
 import FolderComponent from "../components/Folder";
 import MultiSliderComponent from "../components/MultiSlider";
 import UploadButtonComponent from "../components/UploadButton";
+import ProgressBarComponent from "../components/ProgressBar";
 
 /** Root of generated inputs. */
 export default function GeneratedGuiContainer({
@@ -119,6 +120,8 @@ function GeneratedInput(props: { guiId: string }) {
       return <RgbaComponent {...conf} />;
     case "GuiAddButtonGroupMessage":
       return <ButtonGroupComponent {...conf} />;
+    case "GuiAddProgressBarMessage":
+      return <ProgressBarComponent {...conf} />;
     default:
       assertNeverType(conf);
   }