Renamed DownloadableLink as Asset and added to_xml() to it.

examples/simulation/main.py

@@ -5,7 +5,8 @@
 from transformations import compute_motor_torques
 
 from onshape_api.log import LOGGER, LogLevel
-from onshape_api.robot import get_robot
+
+# from onshape_api.robot import RobotType, get_robot
 from onshape_api.utilities import save_gif
 
 HEIGHT = 480
@@ -56,12 +57,6 @@ def get_theta(data):
     return theta[0], theta[1], theta[2]
 
 
-def key_callback(keycode):
-    if chr(keycode) == "e":
-        print("Exiting")
-        return False
-
-
 def control(data, roll_sp=0, pitch_sp=0):
     roll, pitch, yaw = get_theta(data)
     roll = roll - np.pi
@@ -81,35 +76,35 @@ def control(data, roll_sp=0, pitch_sp=0):
     data.ctrl[2] = t2
     data.ctrl[1] = t3
 
-    # print(f"Roll {roll}, Pitch: {pitch}")
+    print(f"Roll {roll}, Pitch: {pitch}")
 
 
 if __name__ == "__main__":
     LOGGER.set_file_name("sim.log")
     LOGGER.set_stream_level(LogLevel.INFO)
 
-    ballbot = get_robot(
-        url="https://cad.onshape.com/documents/1f42f849180e6e5c9abfce52/w/0c00b6520fac5fada24b2104/e/c96b40ef586e60c182f41d29",
-        robot_name="ballbot",
-    )
+    # TODO: Add native support for MJCF (XML) exports: #17
+    # ballbot = get_robot(
+    #     url="https://cad.onshape.com/documents/1f42f849180e6e5c9abfce52/w/0c00b6520fac5fada24b2104/e/c96b40ef586e60c182f41d29",
+    #     robot_name="ballbot",
+    #     robot_type=RobotType.MJCF,
+    # )
+    # ballbot.show()
 
-    ballbot.save()
+    model = mujoco.MjModel.from_xml_path(filename="ballbot.xml")
+    data = mujoco.MjData(model)
 
-    # model = mujoco.MjModel.from_xml_path(filename=urdf_path)
-    # mujoco.mj_saveLastXML("ballbot.xml", model)
-    # data = mujoco.MjData(model)
+    # run_simulation(model, data, 20, 60)
 
-    # # run_simulation(model, data, 20, 60)
+    mujoco.mj_resetData(model, data)
 
-    # mujoco.mj_resetData(model, data)
+    with mujoco.viewer.launch_passive(model, data) as viewer:
+        initial_roll, initial_pitch, initial_yaw = get_theta(data)
 
-    # with mujoco.viewer.launch_passive(model, data, key_callback=key_callback) as viewer:
-    #     initial_roll, initial_pitch, initial_yaw = get_theta(data)
-
-    #     while viewer.is_running():
-    #         mujoco.mj_step(model, data)
-    #         mujoco.mj_forward(model, data)
+        while viewer.is_running():
+            mujoco.mj_step(model, data)
+            mujoco.mj_forward(model, data)
 
-    #         control(data)
+            control(data)
 
-    #         viewer.sync()
+            viewer.sync()

onshape_api/connect.py

@@ -21,6 +21,7 @@
 import secrets
 import string
 import time
+import xml.etree.ElementTree as ET
 from enum import Enum
 from typing import Any, BinaryIO, Optional
 from urllib.parse import parse_qs, urlencode, urlparse
@@ -984,7 +985,7 @@ def _make_headers(self, method, path, query=None, headers=None):
         return req_headers
 
 
-class DownloadableLink:
+class Asset:
     """
     Represents a set of parameters required to download a link from Onshape.
     """
@@ -1002,7 +1003,7 @@ def __init__(
         partID: Optional[str] = None,
     ) -> None:
         """
-        Initialize the DownloadableLink object.
+        Initialize the Asset object.
 
         Args:
             did: The unique identifier of the document.
@@ -1079,3 +1080,26 @@ async def download(self) -> None:
                 LOGGER.info(f"Mesh file saved: {self.absolute_path}")
         except Exception as e:
             LOGGER.error(f"Failed to download {self.file_name}: {e}")
+
+    def to_xml(self, root: ET.Element | None = None) -> str:
+        """
+        Returns the XML representation of the asset, which is a mesh file.
+
+        Examples:
+            >>> asset = Asset(
+            ...     did="a1c1addf75444f54b504f25c",
+            ...     wtype="w",
+            ...     wid="0d17b8ebb2a4c76be9fff3c7",
+            ...     eid="a86aaf34d2f4353288df8812",
+            ...     client=client,
+            ...     transform=np.eye(4),
+            ...     file_name="mesh.stl",
+            ...     is_rigid_assembly=True
+            ... )
+            >>> asset.to_xml()
+            <mesh name="Part-1-1" file="Part-1-1.stl" />
+        """
+        asset = ET.Element("mesh") if root is None else ET.SubElement(root, "mesh")
+        asset.set("mesh", self.file_name)
+        asset.set("file", self.relative_path)
+        return asset

onshape_api/robot.py

@@ -14,7 +14,7 @@
 
 from defusedxml import minidom
 
-from onshape_api.connect import Client, DownloadableLink
+from onshape_api.connect import Asset, Client
 from onshape_api.graph import create_graph
 from onshape_api.log import LOGGER
 from onshape_api.models.document import Document
@@ -65,7 +65,7 @@ def __init__(
         name: str,
         links: dict[str, Link],
         joints: dict[str, BaseJoint],
-        assets: Optional[dict[str, DownloadableLink]] = None,
+        assets: Optional[dict[str, Asset]] = None,
         robot_type: RobotType = RobotType.URDF,
     ):
         self.name = name
@@ -89,12 +89,24 @@ def to_xml(self, robot_type: RobotType) -> ET.Element:
             >>> robot.to_xml()
             <Element 'robot' at 0x7f8b3c0b4c70>
         """
-        robot = ET.Element("robot", name=self.name)
-        for link in self.links.values():
-            link.to_xml(robot)
+        if robot_type == RobotType.URDF:
+            robot = ET.Element("robot", name=self.name)
+            for link in self.links.values():
+                link.to_xml(robot)
+
+            for joint in self.joints.values():
+                joint.to_xml(robot)
+
+        elif robot_type == RobotType.MJCF:
+            robot = ET.Element("mujoco", model=self.name)
+
+            # create an asset element to hold all the assets(meshes)
+            if self.assets:
+                assets_element = ET.SubElement(robot, "asset")
+
+                for asset in self.assets.values():
+                    asset.to_xml(assets_element)
 
-        for joint in self.joints.values():
-            joint.to_xml(robot)
         return robot
 
     def save(self) -> None:
@@ -173,6 +185,7 @@ def get_robot(
     max_traversal_depth: int = 5,
     use_user_defined_root: bool = False,
     save_assembly_as_json: bool = False,
+    robot_type: RobotType = RobotType.URDF,
 ) -> Robot:
     """
     Get a robot model from an Onshape assembly. A convenience function that combines the parsing and
@@ -241,10 +254,11 @@ def get_robot(
         client=client,
     )
 
-    return Robot(name=robot_name, links=links, joints=joints, assets=assets)
+    return Robot(name=robot_name, links=links, joints=joints, assets=assets, robot_type=robot_type)
 
 
 if __name__ == "__main__":
+    LOGGER.set_file_name("robot.log")
     robot = Robot(
         name="Test",
         links={
@@ -254,7 +268,31 @@ def get_robot(
         joints={
             "joint1": FixedJoint(name="joint1", parent="link1", child="link2", origin=Origin.zero_origin()),
         },
+        assets={
+            "link1": Asset(
+                did="1f42f849180e6e5c9abfce52",
+                wtype="w",
+                wid="0c00b6520fac5fada24b2104",
+                eid="c96b40ef586e60c182f41d29",
+                client=Client(env="E:/onshape-api/tests/.env"),
+                transform=(0, 0, 0, 0, 0, 0),
+                file_name="link1.stl",
+                is_rigid_assembly=False,
+                partID="KHD",
+            ),
+            "link2": Asset(
+                did="1f42f849180e6e5c9abfce52",
+                wtype="w",
+                wid="0c00b6520fac5fada24b2104",
+                eid="c96b40ef586e60c182f41d29",
+                client=Client(env="E:/onshape-api/tests/.env"),
+                transform=(0, 0, 0, 0, 0, 0),
+                file_name="link2.stl",
+                is_rigid_assembly=False,
+                partID="KHD",
+            ),
+        },
         robot_type=RobotType.MJCF,
     )
 
-    robot.save()
+    robot.show()

onshape_api/urdf.py

@@ -10,7 +10,7 @@
 import numpy as np
 from networkx import DiGraph
 
-from onshape_api.connect import Client, DownloadableLink
+from onshape_api.connect import Asset, Client
 from onshape_api.log import LOGGER
 from onshape_api.models.assembly import (
     Assembly,
@@ -60,7 +60,7 @@ def get_robot_link(
     wid: str,
     client: Client,
     mate: Optional[Union[MateFeatureData, None]] = None,
-) -> tuple[Link, np.matrix, DownloadableLink]:
+) -> tuple[Link, np.matrix, Asset]:
     """
     Generate a URDF link from an Onshape part.
 
@@ -113,7 +113,7 @@ def get_robot_link(
         wtype = WorkspaceType.W.value
         mvwid = wid
 
-    _downloadable_link = DownloadableLink(
+    _downloadable_link = Asset(
         did=part.documentId,
         wtype=wtype,
         wid=mvwid,
@@ -375,7 +375,7 @@ def get_urdf_components(
     mates: dict[str, MateFeatureData],
     relations: dict[str, MateRelationFeatureData],
     client: Client,
-) -> tuple[dict[str, Link], dict[str, BaseJoint], dict[str, DownloadableLink]]:
+) -> tuple[dict[str, Link], dict[str, BaseJoint], dict[str, Asset]]:
     """
     Generate URDF links and joints from an Onshape assembly.