from __future__ import annotations
import os
import sys
import xml.etree.ElementTree as StdET
from io import BytesIO
from pathlib import Path
if sys.version_info >= (3, 11):
from typing import Self
else:
from typing_extensions import Self
from typing import TYPE_CHECKING
import defusedxml.ElementTree as ET
if TYPE_CHECKING:
from collections.abc import Sequence
from xml.etree.ElementTree import Element, ElementTree
[docs]
class Builder:
"""A class to build and manipulate MuJoCo XML models."""
def __init__(self, *inputs: str, meshdir: str = "meshes/") -> None:
if not inputs:
msg = "Input is required to initialize the Builder"
raise ValueError(msg)
if not all(isinstance(i, str) for i in inputs):
msg = "Input must be an XML string or a file path"
raise TypeError(msg)
self.meshdir = meshdir
self.tree, self.root = self._parse_input(inputs[0])
for other in inputs[1:]:
self += Builder(other, meshdir=meshdir)
[docs]
@staticmethod
def merge(inputs: Sequence[str | Builder], meshdir: str = "meshes/") -> Builder:
"""Merge multiple Builder objects and/or XML strings into one Builder.
Args:
inputs: Sequence of Builder objects and/or XML strings or file paths.
meshdir: Mesh directory (default: "meshes/").
Returns:
Merged Builder instance.
Raises:
ValueError: If no inputs are provided.
"""
if not inputs:
msg = "No inputs provided for merging."
raise ValueError(msg)
builders = [i for i in inputs if isinstance(i, Builder)]
strings = [i for i in inputs if isinstance(i, str)]
if builders:
builder = sum(builders[1:], builders[0])
if strings:
builder += Builder(*strings, meshdir=meshdir)
else:
builder = Builder(*strings, meshdir=meshdir)
return builder
def _parse_input(self, xml_input: str) -> tuple[ElementTree, Element]:
# Parse XML from string or file
if xml_input.strip().startswith("<"):
root = ET.fromstring(xml_input)
else:
path = Path(xml_input)
if not path.exists():
msg = f"File not found: {xml_input}"
raise FileNotFoundError(msg)
root = ET.parse(path).getroot()
# If root is <robot>, ensure <mujoco> child exists (not as wrapper)
if root.tag == "robot":
mujoco_tag = root.find("mujoco")
if mujoco_tag is None:
mujoco_tag = StdET.Element("mujoco")
# Insert <mujoco> as first child (after comments, if any)
insert_idx = 0
for idx, child in enumerate(list(root)):
if not isinstance(child.tag, str) or child.tag.startswith("#"):
insert_idx = idx + 1
else:
break
root.insert(insert_idx, mujoco_tag)
# Ensure <compiler> exists under <mujoco>
compiler_tag = mujoco_tag.find("compiler")
if compiler_tag is None:
compiler_tag = StdET.Element(
"compiler",
{
"angle": "radian",
"meshdir": self.meshdir,
"balanceinertia": "true",
"discardvisual": "true",
},
)
mujoco_tag.insert(0, compiler_tag)
return self._to_safe_tree(root), root
# If root is <mujoco>, ensure <compiler> exists
if root.tag == "mujoco":
compiler_tag = root.find("compiler")
if compiler_tag is None:
compiler_tag = StdET.Element(
"compiler",
{
"angle": "radian",
"meshdir": self.meshdir,
"balanceinertia": "true",
"discardvisual": "true",
},
)
root.insert(0, compiler_tag)
return self._to_safe_tree(root), root
# If root is neither <robot> nor <mujoco>, wrap in <mujoco> and inject <compiler>
mujoco_tag = StdET.Element("mujoco")
mujoco_tag.append(root)
compiler_tag = mujoco_tag.find("compiler")
if compiler_tag is None:
compiler_tag = StdET.Element(
"compiler",
{
"angle": "radian",
"meshdir": self.meshdir,
"balanceinertia": "true",
"discardvisual": "true",
},
)
mujoco_tag.insert(0, compiler_tag)
return self._to_safe_tree(mujoco_tag), mujoco_tag
def _to_safe_tree(self, root: Element) -> ElementTree:
xml_string = StdET.tostring(root)
return ET.parse(BytesIO(xml_string))
def __add__(self, other: str | Builder) -> Self:
if isinstance(other, str):
_, other_root = Builder(other, meshdir=self.meshdir)._parse_input(other)
elif isinstance(other, Builder):
other_root = other.root
else:
msg = "Can only merge with str or Builder"
raise TypeError(msg)
# Determine merge context: MJCF or URDF
if self.root.tag == "robot":
mujoco_self = self.root.find("mujoco")
mujoco_other = (
other_root.find("mujoco")
if other_root.tag == "robot"
else other_root
if other_root.tag == "mujoco"
else None
)
if mujoco_self is not None and mujoco_other is not None:
self._merge_mujoco_tags(mujoco_self, mujoco_other)
elif self.root.tag == "mujoco":
mujoco_self = self.root
mujoco_other = (
other_root.find("mujoco")
if other_root.tag == "robot"
else other_root
if other_root.tag == "mujoco"
else None
)
if mujoco_other is not None:
self._merge_mujoco_tags(mujoco_self, mujoco_other)
else:
# Fallback: merge at root
self._merge_tag("asset", self.root, other_root)
self._merge_tag("worldbody", self.root, other_root)
return self
def _merge_mujoco_tags(self, mujoco_self: Element, mujoco_other: Element) -> None:
# Merge all relevant tags under <mujoco>
for tag in [
"asset",
"worldbody",
"camera",
"light",
"contact",
"equality",
"sensor",
"actuator",
"default",
"tendon",
"include",
]:
self._merge_tag(tag, mujoco_self, mujoco_other)
def _merge_tag(self, tag: str, root1: Element, root2: Element) -> None:
s1, s2 = root1.find(tag), root2.find(tag)
if s1 is None and s2 is not None:
s1 = StdET.SubElement(root1, tag)
if s1 is not None and s2 is not None:
for el in list(s2):
s1.append(el)
[docs]
def save(self, file_path: str) -> str:
if self.tree is not None:
self._indent_xml(self.root)
self.tree.write(file_path, encoding="utf-8", xml_declaration=True)
else:
msg = "No model loaded. Cannot save."
raise ValueError(msg)
return os.path.abspath(file_path)
def _indent_xml(self, elem: Element, level: int = 0) -> None:
i = "\n" + level * " "
if len(elem):
if not elem.text or not elem.text.strip():
elem.text = i + " "
if not elem.tail or not elem.tail.strip():
elem.tail = i
for sub_elem in elem:
self._indent_xml(sub_elem, level + 1)
if not elem[-1].tail or not elem[-1].tail.strip():
elem[-1].tail = i
elif level and (not elem.tail or not elem.tail.strip()):
elem.tail = i
def __str__(self) -> str:
return StdET.tostring(self.root, encoding="unicode", method="xml")
def __repr__(self) -> str:
return self.__str__()
def __len__(self) -> int:
return len(self.root) if self.root is not None else 0
def __radd__(self, other: str | Builder) -> Self:
return self.__add__(other)