lint

Author: WT-MM

urdf2mjcf/convert.py

@@ -533,8 +533,10 @@ def convert_urdf_to_mjcf(
         raise FileNotFoundError(f"URDF file not found: {urdf_path}")
     mjcf_path.parent.mkdir(parents=True, exist_ok=True)
 
-    urdf_tree: ET.ElementTree = ET.parse(urdf_path)
-    robot: ET.Element = urdf_tree.getroot()
+    urdf_tree = ET.parse(urdf_path)
+    robot = urdf_tree.getroot()
+    if robot is None:
+        raise ValueError("URDF file has no root element")
 
     if metadata_file is not None and metadata is not None:
         raise ValueError("Cannot specify both metadata and metadata_file")
@@ -554,10 +556,6 @@ def convert_urdf_to_mjcf(
         joint_metadata, actuator_metadata = _get_empty_joint_and_actuator_metadata(robot)
     assert joint_metadata is not None and actuator_metadata is not None
 
-    # Parse the URDF file.
-    urdf_tree: ET.ElementTree = ET.parse(urdf_path)
-    robot: ET.Element = urdf_tree.getroot()
-
     # Parse materials from URDF - both from root level and from link visuals
     materials: dict[str, str] = {}
 

urdf2mjcf/postprocess/add_sensors.py

@@ -100,17 +100,17 @@ def find_site(body_name: str, site_name: str, site_metadata: list[SiteMetadata])
         # Find or create the site within the body
         site_elem = link_body.find(f"./site[@name='{site_name}']")
         if site_elem is None:
-            site_metadata = next(
+            site_meta = next(
                 (s for s in site_metadata if s.name == site_name), SiteMetadata(body_name=body_name, name=site_name)
             )
 
             site_elem = ET.SubElement(link_body, "site", name=site_name)
-            if site_metadata.size is not None:
-                site_elem.attrib["size"] = " ".join(str(x) for x in site_metadata.size)
-            if site_metadata.pos is not None:
-                site_elem.attrib["pos"] = " ".join(str(x) for x in site_metadata.pos)
-            if site_metadata.site_type is not None:
-                site_elem.attrib["type"] = site_metadata.site_type
+            if site_meta.size is not None:
+                site_elem.attrib["size"] = " ".join(str(x) for x in site_meta.size)
+            if site_meta.pos is not None:
+                site_elem.attrib["pos"] = " ".join(str(x) for x in site_meta.pos)
+            if site_meta.site_type is not None:
+                site_elem.attrib["type"] = site_meta.site_type
             logger.info(f"Created site '{site_name}' on body '{body_name}'.")
 
         return site_elem

urdf2mjcf/utils.py

@@ -1,5 +1,7 @@
 """Defines utility functions."""
 
+from __future__ import annotations
+
 import io
 import re
 import xml.etree.ElementTree as ET
@@ -52,10 +54,18 @@ def iter_meshes(
             yield visual_mesh, collision_mesh
 
 
-def save_xml(path: str | Path | io.StringIO, tree: ET.ElementTree | ET.Element) -> None:
+def save_xml(path: str | Path | io.StringIO, tree: ET.ElementTree[ET.Element] | ET.Element) -> None:
+    """Save XML to file with pretty formatting."""
+    element: ET.Element
     if isinstance(tree, ET.ElementTree):
-        tree = tree.getroot()
-    xmlstr = minidom.parseString(ET.tostring(tree)).toprettyxml(indent="  ")
+        root = tree.getroot()
+        if root is None:
+            raise ValueError("ElementTree has no root element")
+        element = root
+    else:
+        element = tree
+
+    xmlstr = minidom.parseString(ET.tostring(element)).toprettyxml(indent="  ")
     xmlstr = re.sub(r"\n\s*\n", "\n", xmlstr)
 
     # Add newlines between second-level nodes