From 9eedaacf6fbb2e41442abf544405caf493473e39 Mon Sep 17 00:00:00 2001
From: Dominic Reber <dominic@aica.tech>
Date: Sun, 8 Dec 2024 09:32:58 +0000
Subject: [PATCH] feat: use state type

---
 CHANGELOG.md                                  |  1 +
 .../translators/message_writers.py            | 56 ++++++++++---------
 2 files changed, 31 insertions(+), 26 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index dffdade6..7c57d5ba 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -25,6 +25,7 @@ Release Versions:
  - feat(controllers): add TF listener in BaseControllerInterface (#169)
  - feat(controllers): add TF broadcaster in BaseControllerInterface (#170)
  - test(controllers): add TF listener and broadcaster tests (#172)
+ - feat(components): get clproto message type from attribute (#175)
 
 ## 5.0.2
 
diff --git a/source/modulo_core/modulo_core/translators/message_writers.py b/source/modulo_core/modulo_core/translators/message_writers.py
index c228e07a..5cbe4325 100644
--- a/source/modulo_core/modulo_core/translators/message_writers.py
+++ b/source/modulo_core/modulo_core/translators/message_writers.py
@@ -15,42 +15,46 @@
 
 
 def get_clproto_msg_type(state: StateT) -> clproto.MessageType:
-    if isinstance(state, sr.CartesianPose):
+    if not isinstance(state, sr.State) or not hasattr(state, 'get_type') or not callable(state.get_type):
+        return clproto.MessageType.UNKNOWN_MESSAGE
+
+    state_type = state.get_type()
+    if state_type == sr.StateType.STATE:
+        return clproto.MessageType.STATE_MESSAGE
+    elif state_type == sr.StateType.SPATIAL_STATE:
+        return clproto.MessageType.SPATIAL_STATE_MESSAGE
+    elif state_type == sr.StateType.CARTESIAN_STATE:
+        return clproto.MessageType.CARTESIAN_STATE_MESSAGE
+    elif state_type == sr.StateType.CARTESIAN_POSE:
         return clproto.MessageType.CARTESIAN_POSE_MESSAGE
-    elif isinstance(state, sr.CartesianTwist):
+    elif state_type == sr.StateType.CARTESIAN_TWIST:
         return clproto.MessageType.CARTESIAN_TWIST_MESSAGE
-    elif isinstance(state, sr.CartesianAcceleration):
+    elif state_type == sr.StateType.CARTESIAN_ACCELERATION:
         return clproto.MessageType.CARTESIAN_ACCELERATION_MESSAGE
-    elif isinstance(state, sr.CartesianWrench):
+    elif state_type == sr.StateType.CARTESIAN_WRENCH:
         return clproto.MessageType.CARTESIAN_WRENCH_MESSAGE
-    elif isinstance(state, sr.CartesianState):
-        return clproto.MessageType.CARTESIAN_STATE_MESSAGE
-    elif isinstance(state, sr.SpatialState):
-        return clproto.MessageType.SPATIAL_STATE_MESSAGE
-    elif isinstance(state, sr.JointPositions):
+    elif state_type == sr.StateType.JACOBIAN:
+        return clproto.MessageType.JACOBIAN_MESSAGE
+    elif state_type == sr.StateType.JOINT_STATE:
+        return clproto.MessageType.JOINT_STATE_MESSAGE
+    elif state_type == sr.StateType.JOINT_POSITIONS:
         return clproto.MessageType.JOINT_POSITIONS_MESSAGE
-    elif isinstance(state, sr.JointVelocities):
+    elif state_type == sr.StateType.JOINT_VELOCITIES:
         return clproto.MessageType.JOINT_VELOCITIES_MESSAGE
-    elif isinstance(state, sr.JointAccelerations):
-        return clproto.MessageType.JOINT_VELOCITIES_MESSAGE
-    elif isinstance(state, sr.JointTorques):
+    elif state_type == sr.StateType.JOINT_ACCELERATIONS:
+        return clproto.MessageType.JOINT_ACCELERATIONS_MESSAGE
+    elif state_type == sr.StateType.JOINT_TORQUES:
         return clproto.MessageType.JOINT_TORQUES_MESSAGE
-    elif isinstance(state, sr.JointState):
-        return clproto.MessageType.JOINT_STATE_MESSAGE
-    elif isinstance(state, sr.Parameter):
+    elif state_type == sr.StateType.GEOMETRY_SHAPE:
+        return clproto.MessageType.SHAPE_MESSAGE
+    elif state_type == sr.StateType.GEOMETRY_ELLIPSOID:
+        return clproto.MessageType.ELLIPSOID_MESSAGE
+    elif state_type == sr.StateType.PARAMETER:
         return clproto.MessageType.PARAMETER_MESSAGE
-    elif isinstance(state, sr.DigitalIOState):
+    elif state_type == sr.StateType.DIGITAL_IO_STATE:
         return clproto.MessageType.DIGITAL_IO_STATE_MESSAGE
-    elif isinstance(state, sr.AnalogIOState):
+    elif state_type == sr.StateType.ANALOG_IO_STATE:
         return clproto.MessageType.ANALOG_IO_STATE_MESSAGE
-    elif isinstance(state, sr.Jacobian):
-        return clproto.MessageType.JACOBIAN_MESSAGE
-    elif isinstance(state, sr.Ellipsoid):
-        return clproto.MessageType.ELLIPSOID_MESSAGE
-    elif isinstance(state, sr.Shape):
-        return clproto.MessageType.SHAPE_MESSAGE
-    elif isinstance(state, sr.State):
-        return clproto.MessageType.STATE_MESSAGE
     return clproto.MessageType.UNKNOWN_MESSAGE