raise key error if joint name not found (#22)

petrikvladimir · Dec 10, 2024 · fa7c208 · fa7c208
1 parent 57d8005
commit fa7c208
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 0 deletions.
diff --git a/src/robomeshcat/robot.py b/src/robomeshcat/robot.py
@@ -159,6 +159,8 @@ def _get_joint_id(self, key: str | int):
         # index
         if len(self._model.names) == len(self._q) + 1:
             jid -= 1
+        if jid == self._model.nq:
+            raise KeyError(f'Joint {key} not found.')
         return jid
 
     def __getitem__(self, key):

diff --git a/tests/test_robot_q.py b/tests/test_robot_q.py
@@ -18,6 +18,10 @@ def test_q_by_name(self):
         self.assertAlmostEqual(robot["shoulder_pan_joint"], 0.1)
         self.assertAlmostEqual(robot[0], 0.1)
 
+    def test_q_by_name_wrong_key(self):
+        robot = Robot(urdf_path=Path(__file__).parent / 'test_urdf.urdf')
+        self.assertRaises(KeyError, lambda: robot["wrong_key"])
+
 
 if __name__ == '__main__':
     unittest.main()