refactor: use enum for Viser display modes

Author: TomCC7

dimos/manipulation/visualization/viser/gui.py

@@ -14,14 +14,13 @@
 
 from __future__ import annotations
 
-from typing import TypeAlias, cast
+from typing import TypeAlias
 
 from dimos.manipulation.visualization.types import RobotInfo, TargetEvaluation
 from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter
 from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig
 from dimos.manipulation.visualization.viser.runtime import VISER_INSTALL_HINT
 from dimos.manipulation.visualization.viser.scene import (
-    ROBOT_DISPLAY_MODE_VALUES,
     RobotDisplayMode,
     ViserManipulationScene,
 )
@@ -72,8 +71,10 @@
 
 # Fallback joint-slider range (radians) when a robot config omits joint limits.
 DEFAULT_JOINT_LIMITS = (-3.14, 3.14)
-ROBOT_DISPLAY_LABELS = tuple(mode.title() for mode in ROBOT_DISPLAY_MODE_VALUES)
-ROBOT_DISPLAY_MODES = {mode: mode.title() for mode in ROBOT_DISPLAY_MODE_VALUES}
+ROBOT_DISPLAY_LABELS = tuple(mode.value.title() for mode in RobotDisplayMode)
+ROBOT_DISPLAY_MODES: dict[str, RobotDisplayMode] = {
+    mode.value: mode for mode in RobotDisplayMode
+}
 ROBOT_DISPLAY_COLLISION_WARNING = (
     "**Collision meshes unavailable.** Showing visual geometry with collision styling."
 )
@@ -237,10 +238,10 @@ def _robot_display_label(self) -> str:
         scene = self.scene
         if scene is None:
             return ROBOT_DISPLAY_LABELS[0]
-        mode = str(scene.robot_display_mode).lower()
+        mode = scene.robot_display_mode.value
         if mode not in ROBOT_DISPLAY_MODES:
             return ROBOT_DISPLAY_LABELS[0]
-        return ROBOT_DISPLAY_MODES[cast("RobotDisplayMode", mode)]
+        return ROBOT_DISPLAY_MODES[mode].value.title()
 
     def _set_robot_display_mode(self, label: str) -> None:
         if self._closed or self.scene is None:
@@ -249,7 +250,7 @@ def _set_robot_display_mode(self, label: str) -> None:
         mode = str(label).lower()
         if mode not in ROBOT_DISPLAY_MODES:
             return
-        scene.robot_display_mode = cast("RobotDisplayMode", mode)
+        scene.robot_display_mode = ROBOT_DISPLAY_MODES[mode]
         self._sync_robot_display_dropdown()
         self._sync_robot_display_warning()
 
@@ -265,9 +266,9 @@ def _sync_robot_display_warning(self) -> None:
         scene = self.scene
         if handle is None or self._closed or scene is None:
             return
-        mode = str(scene.robot_display_mode).lower()
+        mode = scene.robot_display_mode
         has_collision = scene.collision_geometry_available
-        visible = mode in {"collision", "both"} and not has_collision
+        visible = mode in {RobotDisplayMode.COLLISION, RobotDisplayMode.BOTH} and not has_collision
         self._set_optional_handle_attr(handle, "visible", visible)
 
     def _set_scene_grid_visible(self, visible: bool) -> None:

dimos/manipulation/visualization/viser/scene.py

@@ -15,8 +15,9 @@
 from __future__ import annotations
 
 from collections.abc import Callable, Sequence
+from enum import StrEnum
 from pathlib import Path
-from typing import Literal, Protocol, TypeAlias, cast, get_args
+from typing import Protocol, TypeAlias, cast
 
 from yourdfpy import URDF  # type: ignore[import-untyped]
 
@@ -73,11 +74,10 @@
 COLLISION_MESH_COLOR = (210, 40, 220)
 COLLISION_MESH_OPACITY = 0.35
 
-RobotDisplayMode: TypeAlias = Literal["visual", "collision", "both"]
-ROBOT_DISPLAY_MODE_VALUES: tuple[RobotDisplayMode, ...] = get_args(RobotDisplayMode)
-_VISUAL_DISPLAY_MODE = ROBOT_DISPLAY_MODE_VALUES[0]
-_COLLISION_DISPLAY_MODE = ROBOT_DISPLAY_MODE_VALUES[1]
-_BOTH_DISPLAY_MODE = ROBOT_DISPLAY_MODE_VALUES[2]
+class RobotDisplayMode(StrEnum):
+    VISUAL = "visual"
+    COLLISION = "collision"
+    BOTH = "both"
 
 SceneHandle: TypeAlias = ViserUrdf | TransformControlsHandle | GridHandle | MeshHandle
 
@@ -105,7 +105,7 @@ def __init__(
         self._preview_visible: dict[str, bool] = {}
         self._target_tracks_current: dict[str, bool] = {}
         self._collision_fallback_urdfs: dict[str, ViserUrdf] = {}
-        self._robot_display_mode: RobotDisplayMode = "visual"
+        self._robot_display_mode = RobotDisplayMode.VISUAL
         self._ensure_reference_grid()
 
     @property
@@ -114,11 +114,13 @@ def robot_display_mode(self) -> RobotDisplayMode:
         return self._robot_display_mode
 
     @robot_display_mode.setter
-    def robot_display_mode(self, mode: RobotDisplayMode) -> None:
+    def robot_display_mode(self, mode: RobotDisplayMode | str) -> None:
         """Set the primary robot display mode and apply it immediately."""
-        if mode not in ROBOT_DISPLAY_MODE_VALUES:
-            raise ValueError(f"Unsupported robot display mode: {mode!r}")
-        self._robot_display_mode = mode
+        try:
+            normalized_mode = RobotDisplayMode(mode)
+        except ValueError as error:
+            raise ValueError(f"Unsupported robot display mode: {mode!r}") from error
+        self._robot_display_mode = normalized_mode
         for robot_id in self._configs_by_id:
             self._apply_robot_display_mode(robot_id)
 
@@ -319,7 +321,7 @@ def close(self) -> None:
         self._configs_by_id.clear()
         self._preview_visible.clear()
         self._target_tracks_current.clear()
-        self._robot_display_mode = "visual"
+        self._robot_display_mode = RobotDisplayMode.VISUAL
 
     def _ensure_robot_urdfs(self, robot_id: str, config: RobotModelConfig) -> None:
         if not config.model_path:
@@ -416,17 +418,21 @@ def _apply_robot_display_mode(self, robot_id: str) -> None:
         # Viser's public flags manage all links, including links whose mesh
         # handles are not exposed by the helper.  A model without collision
         # geometry falls back to its visual representation.
-        current.show_visual = mode in {_VISUAL_DISPLAY_MODE, _BOTH_DISPLAY_MODE}
+        current.show_visual = mode in {RobotDisplayMode.VISUAL, RobotDisplayMode.BOTH}
         current.show_collision = has_collision and mode in {
-            _COLLISION_DISPLAY_MODE,
-            _BOTH_DISPLAY_MODE,
+            RobotDisplayMode.COLLISION,
+            RobotDisplayMode.BOTH,
         }
         fallback = self._collision_fallback_urdfs.get(robot_id)
         if fallback is not None:
-            fallback.show_visual = mode in {_COLLISION_DISPLAY_MODE, _BOTH_DISPLAY_MODE}
+            fallback.show_visual = mode in {
+                RobotDisplayMode.COLLISION,
+                RobotDisplayMode.BOTH,
+            }
             fallback.show_collision = False
             self._set_handle_visibility(
-                fallback, mode in {_COLLISION_DISPLAY_MODE, _BOTH_DISPLAY_MODE}
+                fallback,
+                mode in {RobotDisplayMode.COLLISION, RobotDisplayMode.BOTH},
             )
 
     def prepared_urdf_path(self, config: RobotModelConfig) -> Path:

dimos/manipulation/visualization/viser/test_gui_status.py

@@ -26,7 +26,7 @@
 from dimos.manipulation.visualization.viser.adapter import InProcessViserAdapter
 from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig
 from dimos.manipulation.visualization.viser.gui import ROBOT_DISPLAY_LABELS, ViserPanelGui
-from dimos.manipulation.visualization.viser.scene import ViserManipulationScene
+from dimos.manipulation.visualization.viser.scene import RobotDisplayMode, ViserManipulationScene
 from dimos.manipulation.visualization.viser.state import FeasibilityStatus
 
 
@@ -86,19 +86,20 @@ def add_markdown(self, value: str, *, visible: bool = True) -> DisplayMarkdown:
 
 
 class DisplayScene:
-    def __init__(self, mode: str = "visual", has_collision: bool = True) -> None:
-        self._robot_display_mode = mode
+    def __init__(self, mode: RobotDisplayMode | str = "visual", has_collision: bool = True) -> None:
+        self._robot_display_mode = RobotDisplayMode(mode)
         self.collision_geometry_available = has_collision
         self.set_modes: list[str] = []
 
     @property
-    def robot_display_mode(self) -> str:
+    def robot_display_mode(self) -> RobotDisplayMode:
         return self._robot_display_mode
 
     @robot_display_mode.setter
-    def robot_display_mode(self, mode: str) -> None:
-        self.set_modes.append(mode)
-        self._robot_display_mode = mode
+    def robot_display_mode(self, mode: RobotDisplayMode | str) -> None:
+        normalized_mode = RobotDisplayMode(mode)
+        self.set_modes.append(normalized_mode)
+        self._robot_display_mode = normalized_mode
 
     def has_reference_grid(self) -> bool:
         return False

dimos/manipulation/visualization/viser/test_viser_visualization.py

@@ -35,7 +35,7 @@
 )
 from dimos.manipulation.visualization.viser.config import ViserVisualizationConfig
 from dimos.manipulation.visualization.viser.gui import ViserPanelGui
-from dimos.manipulation.visualization.viser.scene import ViserManipulationScene
+from dimos.manipulation.visualization.viser.scene import RobotDisplayMode, ViserManipulationScene
 from dimos.manipulation.visualization.viser.state import (
     ActionStatus,
     FeasibilityStatus,
@@ -1483,7 +1483,7 @@ def test_gui_moves_joint_target_immediately_and_stores_evaluated_joint_solution(
     target_updates = []
     target_pose_updates = []
     scene = SimpleNamespace(
-        robot_display_mode="visual",
+        robot_display_mode=RobotDisplayMode.VISUAL,
         collision_geometry_available=False,
         has_reference_grid=lambda: False,
         ensure_target_controls=lambda *args: None,
@@ -1550,7 +1550,7 @@ def test_gui_cartesian_ik_result_does_not_rewrite_active_gizmo(
     target_joint_updates = []
     target_pose_updates = []
     scene = SimpleNamespace(
-        robot_display_mode="visual",
+        robot_display_mode=RobotDisplayMode.VISUAL,
         collision_geometry_available=False,
         has_reference_grid=lambda: False,
         ensure_target_controls=lambda *args: None,
@@ -1600,7 +1600,7 @@ def test_gui_collision_evaluation_marks_target_infeasible_and_colors_scene(
     adapter = InProcessViserAdapter(world_monitor=world_monitor, manipulation_module=module)
     visual_states = []
     scene = SimpleNamespace(
-        robot_display_mode="visual",
+        robot_display_mode=RobotDisplayMode.VISUAL,
         collision_geometry_available=False,
         has_reference_grid=lambda: False,
         ensure_target_controls=lambda *args: None,