From 9ea25a60d6bc18a76ac3bd3e9b27b98367921e2c Mon Sep 17 00:00:00 2001
From: Guilhem Saurel <guilhem.saurel@laas.fr>
Date: Mon, 16 Sep 2024 17:36:25 +0200
Subject: [PATCH] configurable default visualizer

---
 CHANGELOG.md                                  |  4 ++
 bindings/python/CMakeLists.txt                |  3 +-
 bindings/python/pinocchio/robot_wrapper.py    |  4 +-
 .../python/pinocchio/visualize/__init__.py    |  1 +
 .../python/pinocchio/visualize/visualizers.py | 55 +++++++++++++++++++
 5 files changed, 64 insertions(+), 3 deletions(-)
 create mode 100644 bindings/python/pinocchio/visualize/visualizers.py

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e53b600890..8572cadf9f 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 
+### Added
+
+- Default visualizer can be changed with `PINOCCHIO_VIEWER` environment variable
+
 ### Fixed
 - Fix linkage of Boost.Serialization on Windows ([#2400](https://github.com/stack-of-tasks/pinocchio/pull/2400))
 - Fix mjcf parser appending of inertias at root joint ([#2403](https://github.com/stack-of-tasks/pinocchio/pull/2403))
diff --git a/bindings/python/CMakeLists.txt b/bindings/python/CMakeLists.txt
index 96f5b77af6..b181313f23 100644
--- a/bindings/python/CMakeLists.txt
+++ b/bindings/python/CMakeLists.txt
@@ -225,7 +225,8 @@ if(BUILD_PYTHON_INTERFACE)
   # --- INSTALL VISUALIZATION SCRIPTS
   install_python_files(
     MODULE visualize FILES __init__.py base_visualizer.py gepetto_visualizer.py
-                           meshcat_visualizer.py panda3d_visualizer.py rviz_visualizer.py)
+                           meshcat_visualizer.py panda3d_visualizer.py rviz_visualizer.py
+                           visualizers.py)
 
   # --- STUBS --- #
   if(GENERATE_PYTHON_STUBS)
diff --git a/bindings/python/pinocchio/robot_wrapper.py b/bindings/python/pinocchio/robot_wrapper.py
index fd8ec06082..4749049ee5 100644
--- a/bindings/python/pinocchio/robot_wrapper.py
+++ b/bindings/python/pinocchio/robot_wrapper.py
@@ -413,14 +413,14 @@ def initViewer(self, share_data=True, *args, **kwargs):
         """Init the viewer"""
         # Set viewer to use to MeshCat.
         if self.viz is None:
-            from .visualize import MeshcatVisualizer
+            from .visualize import Visualizer
 
             data, collision_data, visual_data = None, None, None
             if share_data:
                 data = self.data
                 collision_data = self.collision_data
                 visual_data = self.visual_data
-            self.viz = MeshcatVisualizer(
+            self.viz = Visualizer.default()(
                 self.model,
                 self.collision_model,
                 self.visual_model,
diff --git a/bindings/python/pinocchio/visualize/__init__.py b/bindings/python/pinocchio/visualize/__init__.py
index 34ecd8e170..820588cd61 100644
--- a/bindings/python/pinocchio/visualize/__init__.py
+++ b/bindings/python/pinocchio/visualize/__init__.py
@@ -4,3 +4,4 @@
 from .meshcat_visualizer import MeshcatVisualizer
 from .panda3d_visualizer import Panda3dVisualizer
 from .rviz_visualizer import RVizVisualizer
+from .visualizers import Visualizer
diff --git a/bindings/python/pinocchio/visualize/visualizers.py b/bindings/python/pinocchio/visualize/visualizers.py
new file mode 100644
index 0000000000..e1e0002a8b
--- /dev/null
+++ b/bindings/python/pinocchio/visualize/visualizers.py
@@ -0,0 +1,55 @@
+from enum import Enum
+from importlib.util import find_spec
+from os import environ
+
+from .base_visualizer import BaseVisualizer
+from .gepetto_visualizer import GepettoVisualizer
+from .meshcat_visualizer import MeshcatVisualizer
+from .panda3d_visualizer import Panda3dVisualizer
+from .rviz_visualizer import RVizVisualizer
+
+
+class Visualizer(Enum):
+    BASE = BaseVisualizer
+    GEPETTO = GepettoVisualizer
+    MESHCAT = MeshcatVisualizer
+    PANDA3D = Panda3dVisualizer
+    RVIZ = RVizVisualizer
+
+    @classmethod
+    def default(cls):
+        """
+        Allow user to choose their prefered viewer with eg.
+        export PINOCCHIO_VIEWER=RVIZ.
+
+        Otherwise, try to find one which is installed.
+        """
+        # Allow user to define which viewer they want
+        if "PINOCCHIO_VIEWER" in environ:
+            selected = environ["PINOCCHIO_VIEWER"].upper()
+            if hasattr(cls, selected):
+                return getattr(cls, selected).value
+            err = (
+                f"The visualizer '{selected}' is not available.\n"
+                "Please set PINOCCHIO_VIEWER to something installed among:\n"
+                "- meshcat\n"
+                "- gepetto-viewer\n"
+                "- panda3d\n"
+                "- rviz\n"
+            )
+            raise ImportError(err)
+
+        # Otherwise, use the first available
+        for v in ["meshcat", "gepetto", "panda3d_viewer", "rviz"]:
+            if find_spec(v) is not None:
+                return getattr(cls, v.replace("_viewer", "").upper()).value
+
+        err = (
+            "No visualizer could be found.\n"
+            "Please install one of those:\n"
+            "- meshcat\n"
+            "- gepetto-viewer\n"
+            "- panda3d\n"
+            "- rviz\n"
+        )
+        raise ImportError(err)