Add pybind wrappers for BodyNode's getChildBodyNode() & getChil… (#1387)

dartsim · Aug 12, 2019 · 15b5af8 · 15b5af8
1 parent 775a7d3
commit 15b5af8
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 1 deletion.
diff --git a/python/dartpy/dynamics/BodyNode.cpp b/python/dartpy/dynamics/BodyNode.cpp
@@ -866,11 +866,21 @@ void BodyNode(py::module& m)
           +[](const dart::dynamics::BodyNode* self) -> std::size_t {
             return self->getNumChildBodyNodes();
           })
+      .def(
+          "getChildBodyNode",
+          +[](dart::dynamics::BodyNode* self, std::size_t index) -> dart::dynamics::BodyNode* { return self->getChildBodyNode(index); },
+          ::py::arg("index"),
+          ::py::return_value_policy::reference_internal)
       .def(
           "getNumChildJoints",
           +[](const dart::dynamics::BodyNode* self) -> std::size_t {
             return self->getNumChildJoints();
           })
+      .def(
+          "getChildJoint",
+          +[](dart::dynamics::BodyNode* self, std::size_t index) -> dart::dynamics::Joint* { return self->getChildJoint(index); },
+          ::py::arg("index"),
+          ::py::return_value_policy::reference_internal)
       .def(
           "getNumShapeNodes",
           +[](const dart::dynamics::BodyNode* self) -> std::size_t {

diff --git a/python/tests/unit/dynamics/test_body_node.py b/python/tests/unit/dynamics/test_body_node.py
@@ -11,7 +11,8 @@ def test_basic():
 
     for i in range(kr5.getNumBodyNodes()):
         body = kr5.getBodyNode(i)
-        assert np.array_equal(np.array(body.getSpatialVelocity()), np.zeros(6)) is True
+        assert np.array_equal(
+            np.array(body.getSpatialVelocity()), np.zeros(6)) is True
         shape_nodes = body.getShapeNodes()
         for shape_node in shape_nodes:
             print(shape_node)
@@ -24,5 +25,25 @@ def test_basic():
                 dynamics = shape_node.getDynamicsAspect()
 
 
+def testGetChildPMethods():
+    urdfParser = dart.utils.DartLoader()
+    kr5 = urdfParser.parseSkeleton("dart://sample/urdf/KR5/KR5 sixx R650.urdf")
+    assert kr5 is not None
+
+    currentBodyNode = kr5.getRootBodyNode()
+    assert currentBodyNode is not None
+
+    for i in range(1, kr5.getNumBodyNodes()):
+        childBodyNode = currentBodyNode.getChildBodyNode(0)
+        childJoint = currentBodyNode.getChildJoint(0)
+
+        assert childBodyNode is not None
+        assert childJoint is not None
+        assert childBodyNode.getName() == kr5.getBodyNode(i).getName()
+        assert childJoint.getName() == kr5.getJoint(i).getName()
+
+        currentBodyNode = childBodyNode
+
+
 if __name__ == "__main__":
     pytest.main()