def traverse()

in robogym/mujoco/forward_kinematics.py [0:0]

• 27 lines of code
• 6 McCabe index (conditional complexity)

        def traverse(rt: et.Element, parent_body: str):
            assert rt.tag == "body", "only start from body tag in xml"
            matrix = get_matrix(rt)
            name = rt.attrib.get("name", "noname_body_%d" % len(body_info))
            body_info[name] = (matrix, parent_body)

            for x in rt.findall("joint"):
                joint_name = x.attrib.get("name", "")
                joint_idx: int = joint_names_idx.get(joint_name, -1)
                if joint_idx == -1:
                    continue

                assert (
                    x.attrib.get("type", "hinge") == "hinge"
                ), "currently only support hinge joints"

                pos = np.fromstring(x.attrib.get("pos"), sep=" ")
                axis = np.fromstring(x.attrib.get("axis"), sep=" ")

                joint_info[joint_idx] = (pos, axis)

                assert (
                    joint_name not in body_joints
                ), "Only support open chain system, unsupported rigid bodies"
                body_joints[name] = joint_name

            for x in rt.findall("site"):
                site_idx = target_sites_idx.get(x.attrib.get("name", ""), -1)
                if site_idx != -1:
                    matrix = get_matrix(x)
                    site_info[site_idx] = (matrix, name)

            for x in rt.findall("body"):
                # recursive scan through body parts
                traverse(x, name)