From fa7c208d0f46134ab3c5e87ac849e24efb8744e2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Vladim=C3=ADr=20Petr=C3=ADk?= Date: Tue, 10 Dec 2024 14:13:49 +0100 Subject: [PATCH] raise key error if joint name not found (#22) --- src/robomeshcat/robot.py | 2 ++ tests/test_robot_q.py | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/src/robomeshcat/robot.py b/src/robomeshcat/robot.py index 855294e..a16c24e 100644 --- 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 index 18d4656..13b5fc5 100644 --- 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()