Source code for rofunc.utils.robolab.formatter.mjcf_parser.traversal_utils

# Copyright 2018 The dm_control Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================

"""Utility functions that operate on MJCF elements."""

_ACTUATOR_TAGS = ('general', 'motor', 'position',
                  'velocity', 'cylinder', 'muscle')


def get_freejoint(element):
    """Retrieves the free joint of a body. Returns `None` if there isn't one."""
    if element.tag != 'body':
        return None
    elif hasattr(element, 'freejoint') and element.freejoint is not None:
        return element.freejoint
    else:
        joints = element.find_all('joint', immediate_children_only=True)
        for joint in joints:
            if joint.type == 'free':
                return joint
        return None


def get_attachment_frame(mjcf_model):
    return mjcf_model.parent_model.find('attachment_frame', mjcf_model.model)


def get_frame_freejoint(mjcf_model):
    frame = get_attachment_frame(mjcf_model)
    return get_freejoint(frame)


def get_frame_joints(mjcf_model):
    """Retrieves all joints belonging to the attachment frame of an MJCF model."""
    frame = get_attachment_frame(mjcf_model)
    if frame:
        return frame.find_all('joint', immediate_children_only=True)
    else:
        return None


def commit_defaults(element, attributes=None):
    """Commits default values into attributes of the specified element.

    Args:
      element: A PyMJCF element.
      attributes: (optional) A list of strings specifying the attributes to be
        copied from defaults, or `None` if all attributes should be copied.
    """
    dclass = element.dclass
    parent = element.parent
    while dclass is None and parent != element.root:
        dclass = getattr(parent, 'childclass', None)
        parent = parent.parent
    if dclass is None:
        dclass = element.root.default

    while dclass != element.root:
        if element.tag in _ACTUATOR_TAGS:
            tags = _ACTUATOR_TAGS
        else:
            tags = (element.tag,)
        for tag in tags:
            default_element = getattr(dclass, tag)
            for name, value in default_element.get_attributes().items():
                if attributes is None or name in attributes:
                    if hasattr(element, name) and getattr(element, name) is None:
                        setattr(element, name, value)
        dclass = dclass.parent