Add sibling traversal for tree elements.

PiperOrigin-RevId: 691057957
Change-Id: I3f30cf175b5ad0842f3fbbf7fdb42ba7a2212682
This commit is contained in:
Alessio Quaglino
2024-10-29 09:54:07 -07:00
committed by Copybara-Service
parent b941e994d9
commit e835f139c8
2 changed files with 68 additions and 7 deletions
+21 -2
View File
@@ -602,8 +602,12 @@ class SpecsTest(absltest.TestCase):
<worldbody>
<body name="body1">
<body name="body3">
<site name="site1"/>
<site name="site2"/>
<site name="site3"/>
<site name="site4"/>
<body name="body4">
<site name="site"/>
<site name="site5"/>
</body>
</body>
</body>
@@ -614,6 +618,9 @@ class SpecsTest(absltest.TestCase):
spec = mujoco.MjSpec.from_string(main_xml)
bodytype = mujoco.mjtObj.mjOBJ_BODY
self.assertLen(spec.bodies, 5)
self.assertLen(spec.sites, 5)
self.assertLen(spec.worldbody.find_all('body'), 4)
self.assertLen(spec.worldbody.find_all('site'), 5)
self.assertEqual(spec.bodies[1].name, 'body1')
self.assertEqual(spec.bodies[2].name, 'body2')
self.assertEqual(spec.bodies[3].name, 'body3')
@@ -630,7 +637,19 @@ class SpecsTest(absltest.TestCase):
self.assertEqual(spec.bodies[3].find_all('body')[0].name, 'body4')
self.assertEmpty(spec.bodies[2].find_all('body'))
self.assertEmpty(spec.bodies[4].find_all('body'))
self.assertEqual(spec.worldbody.find_all('site')[0].name, 'site')
self.assertEqual(spec.worldbody.find_all('site')[0].name, 'site1')
self.assertEqual(spec.worldbody.find_all('site')[1].name, 'site2')
self.assertEqual(spec.worldbody.find_all('site')[2].name, 'site3')
self.assertEqual(spec.worldbody.find_all('site')[3].name, 'site4')
self.assertEqual(spec.worldbody.find_all('site')[4].name, 'site5')
self.assertEmpty(spec.bodies[2].sites)
self.assertLen(spec.bodies[3].sites, 4)
self.assertLen(spec.bodies[4].sites, 1)
self.assertEqual(spec.bodies[3].sites[0].name, 'site1')
self.assertEqual(spec.bodies[3].sites[1].name, 'site2')
self.assertEqual(spec.bodies[3].sites[2].name, 'site3')
self.assertEqual(spec.bodies[3].sites[3].name, 'site4')
self.assertEqual(spec.bodies[4].sites[0].name, 'site5')
with self.assertRaises(ValueError) as cm:
spec.worldbody.find_all('actuator')
self.assertEqual(