from collections import OrderedDict from decimal import getcontext from os.path import abspath, dirname, join, exists from mujoco_worldgen.transforms import closure_transform import numpy as np import xmltodict import os from mujoco_worldgen.util.types import accepts, returns from mujoco_worldgen.util.path import worldgen_path from mujoco_worldgen.parser.normalize import normalize, stringify getcontext().prec = 4 ''' This directory should contain all XML string processing. No other files should be manually converting types for XML processing. API: parse_file() - takes in path to mujoco file and returns normalized dictionary unparse_dict() - takes in an xml dictionary and returns an XML string NOTE FOR TRANSFORMS: The internal xml_dict layout passed into transforms is the one returned by normalize() -- see that docstring for more details on its layout. Every other method should be considered internal! ''' @accepts(str, bool) @returns(OrderedDict) def parse_file(xml_path, enforce_validation=True): ''' Reads xml from xml_path, consolidates all includes in xml, and returns a normalized xml dictionary. See preprocess() ''' # TODO: use XSS or DTD checking to verify XML structure with open(xml_path) as f: xml_string = f.read() xml_doc_dict = xmltodict.parse(xml_string.strip()) assert 'mujoco' in xml_doc_dict, "XML must contain node" xml_dict = xml_doc_dict['mujoco'] assert isinstance(xml_dict, OrderedDict), \ "Invalid node type {}".format(type(xml_dict)) preprocess(xml_dict, xml_path, enforce_validation=enforce_validation) return xml_dict @accepts(OrderedDict) @returns(str) def unparse_dict(xml_dict): ''' Convert a normalized XML dictionary into a XML string. See stringify(). Note: this modifies xml_dict in place to have strings instead of values. ''' stringify(xml_dict) xml_doc_dict = OrderedDict(mujoco=xml_dict) return xmltodict.unparse(xml_doc_dict, pretty=True) @accepts(OrderedDict, str, bool) def preprocess(xml_dict, root_xml_path, enforce_validation=True): ''' All the steps to turn XML into Worldgen readable form: - normalize: changes strings to floats / vectors / bools, and turns consistently nodes to OrderedDict and List - name_meshes: some meshes are missing names. Here we give default names. - rename_defaults: some defaults are global, we give them names so they won't be anymore. - extract_includes: recursively, we extract includes and merge them. - validate: we apply few final checks on the structure. ''' normalize(xml_dict) set_absolute_paths(xml_dict, root_xml_path) extract_includes(xml_dict, root_xml_path, enforce_validation=enforce_validation) if enforce_validation: validate(xml_dict) @accepts(OrderedDict, str) def set_absolute_paths(xml_dict, root_xml_path): dirnames = ["@meshdir", "@texturedir"] if "compiler" in xml_dict: for drname in dirnames: if drname in xml_dict["compiler"]: asset_dir = worldgen_path('assets') + '/' path = xml_dict["compiler"][drname] if path[0] != "/": relative_path = os.path.dirname(root_xml_path) + "/" + path xml_dict["compiler"][drname] = os.path.abspath(relative_path) elif path.find(asset_dir) > -1: xml_dict["compiler"][drname] = worldgen_path( 'assets', path.split(asset_dir)[-1]) @accepts(OrderedDict, str, bool) def extract_includes(xml_dict, root_xml_path, enforce_validation=True): ''' extracts "include" xmls and substitutes them. ''' def transform_include(node): if "include" in node: if isinstance(node["include"], OrderedDict): node["include"] = [node["include"]] include_xmls = [] for include_dict in node["include"]: include_path = include_dict["@file"] if not exists(include_path): include_path = join(dirname(abspath(root_xml_path)), include_path) assert exists(include_path), "Cannot include file: %s" % include_path with open(include_path) as f: include_string = f.read() include_xml = xmltodict.parse(include_string.strip()) closure_transform(transform_include)(include_xml) assert "mujocoinclude" in include_xml, "Missing ." include_xmls.append(include_xml["mujocoinclude"]) del node["include"] for include_xml in include_xmls: preprocess(include_xml, root_xml_path, enforce_validation=enforce_validation) update_mujoco_dict(node, include_xml) closure_transform(transform_include)(xml_dict) @accepts(OrderedDict, OrderedDict) @returns(None.__class__) def update_mujoco_dict(dict_a, dict_b): ''' Update mujoco dict_a with the contents of another mujoco dict_b. ''' other = (str, int, float, np.ndarray, tuple) for key, value in dict_b.items(): if key not in dict_a: dict_a[key] = value elif isinstance(dict_a[key], list): assert isinstance(value, list), "Expected %s to be a list" % value dict_a[key] += value elif isinstance(value, other): assert(isinstance(dict_a[key], other)) assert dict_a[key] == value, "key=%s\n,Trying to merge dictionaries. " \ "They don't agree on value: %s vs %s" % (key, dict_a[key], value) else: assert isinstance(dict_a[key], OrderedDict), "dict_a = %s\nkey=%s\nExpected dict_a[key] to be a OrderedDict." % (dict_a, key) assert(isinstance(value, OrderedDict)) update_mujoco_dict(dict_a[key], value) @accepts(OrderedDict) def validate(xml_dict): ''' If we make assumptions elsewhere in XML processing, then they should be enforced here. ''' # Assumption: radians for angles, "xyz" euler angle sequence, etc. values = {'@coordinate': 'local', '@angle': 'radian', '@eulerseq': 'xyz'} for key, value in values.items(): if key in xml_dict: assert value == xml_dict[key], 'Invalid value for \"%s\". We support only \"%s\"' % (key, value) # Assumption: all meshes have name if "asset" in xml_dict and "mesh" in xml_dict["asset"]: for mesh in xml_dict["asset"]["mesh"]: assert "@name" in mesh, "%s is missing name" % mesh # Assumption: none all the default classes is global. if "default" in xml_dict: for key, value in xml_dict["default"].items(): assert key == "default", "Dont use global variables in default %s %s" % (key, value) # Assumption: all joints have name. def assert_joint_names(node): if "joint" in node: for joint in node["joint"]: assert "@name" in joint, "Missing name for %s" % joint if "worldbody" in xml_dict: closure_transform(assert_joint_names)(xml_dict["worldbody"])