Added support for assembly names.

benchmark/assembly.py

@@ -3,87 +3,87 @@
 import os
 import pstats
 
+import pandas as pd
+
 from onshape_api.connect import Client
 from onshape_api.graph import create_graph, save_graph
-from onshape_api.models.assembly import Assembly
-from onshape_api.models.document import Document, DocumentMetaData
 from onshape_api.models.robot import Robot
 from onshape_api.parse import (
     get_instances,
-    get_mates,
+    get_mates_and_relations,
     get_occurences,
     get_parts,
     get_subassemblies,
 )
 from onshape_api.urdf import get_urdf_components
 from onshape_api.utilities import LOGGER
-from onshape_api.utilities.helpers import get_random_files, get_sanitized_name
 
 SCRIPT_DIRECTORY = os.path.dirname(os.path.realpath(__file__))
 DATA_DIRECTORY = "/../onshape_api/data"
+ERRORED_ASSEMBLY = "errored_assembly.json"
 
 
-def main(checkpoint: int = 0):
-    client = Client()
-
-    if os.path.exists(f"checkpoint_document_{checkpoint}.json") and os.path.exists(
-        f"checkpoint_assembly_{checkpoint}.json"
-    ):
-        with open(f"checkpoint_document_{checkpoint}.json") as f:
-            document_meta_data = DocumentMetaData.model_validate_json(json.load(f))
-        assembly = Assembly.model_validate_json(json.load(open(f"checkpoint_assembly_{checkpoint}.json")))
+def get_random_assembly(assembly_df: pd.DataFrame) -> dict:
+    return assembly_df.sample().to_dict(orient="records")[0]
 
-    else:
-        json_path = SCRIPT_DIRECTORY + DATA_DIRECTORY + f"/assemblies_checkpoint_{checkpoint}_json"
-        json_file_path, document_id = get_random_files(directory=json_path, file_extension=".json", count=1)
 
-        json_data = json.load(open(json_file_path[0]))
+def get_random_urdf(data_path: str, client: Client):
+    assembly_df = pd.read_parquet(data_path, engine="pyarrow")
 
-        did = document_id[0].split("_")[0]
-        eid = document_id[0].split("_")[1]
+    if os.path.exists(ERRORED_ASSEMBLY):
+        assembly_dict = json.load(open(ERRORED_ASSEMBLY))
+    else:
+        assembly_dict = get_random_assembly(assembly_df)
 
-        document_meta_data = client.get_document_metadata(did)
-        document_meta_data.name = get_sanitized_name(document_meta_data.name)
+    assembly, _ = client.get_assembly(
+        did=assembly_dict["documentId"],
+        wtype=assembly_dict["wtype"],
+        wid=assembly_dict["workspaceId"],
+        eid=assembly_dict["elementId"],
+        with_meta_data=True,
+    )
 
-        document = Document(did=did, wtype="w", wid=document_meta_data.defaultWorkspace.id, eid=eid)
-        assembly = Assembly(**json_data)
-        assembly.document = document
+    assembly_robot_name = f"{assembly.document.name + '-' + assembly.name}"
 
     try:
         instances, id_to_name_map = get_instances(assembly)
         occurences = get_occurences(assembly, id_to_name_map)
         parts = get_parts(assembly, client, instances)
         subassemblies = get_subassemblies(assembly, instances)
-        mates = get_mates(assembly, subassemblies, id_to_name_map)
+        mates, relations = get_mates_and_relations(assembly, subassemblies, id_to_name_map)
 
         graph, root_node = create_graph(occurences=occurences, instances=instances, parts=parts, mates=mates)
-        save_graph(graph, f"{document_meta_data.name}.png")
+        save_graph(graph, f"{assembly_robot_name}.png")
 
         links, joints = get_urdf_components(
-            assembly=assembly, graph=graph, root_node=root_node, parts=parts, mates=mates, client=client
+            assembly=assembly,
+            graph=graph,
+            root_node=root_node,
+            parts=parts,
+            mates=mates,
+            relations=relations,
+            client=client,
         )
 
-        robot = Robot(name=f"{document_meta_data.name}", links=links, joints=joints)
-        robot.save(f"{document_meta_data.name}.urdf")
+        robot = Robot(name=f"{assembly_robot_name}", links=links, joints=joints)
+        robot.save(f"{assembly_robot_name}.urdf")
         LOGGER.info(f"Onshape document: {assembly.document.url}")
-        LOGGER.info(f"URDF saved to {os.path.abspath(f"{document_meta_data.name}.urdf")}")
+        LOGGER.info(f"URDF saved to {os.path.abspath(f"{assembly_robot_name}.urdf")}")
 
     except Exception as e:
-        LOGGER.warning(f"Error processing assembly: {document_meta_data.name}")
+        LOGGER.warning(f"Error processing robot: {assembly_robot_name}")
         LOGGER.warning(e)
         LOGGER.warning(f"Onshape document: {assembly.document.url}")
 
-        with open(f"checkpoint_document_{checkpoint}.json", "w") as f:
-            json.dump(document_meta_data.model_dump_json(), f)
-
-        with open(f"checkpoint_assembly_{checkpoint}.json", "w") as f:
-            json.dump(assembly.model_dump_json(), f)
+        with open(ERRORED_ASSEMBLY, "w") as f:
+            json.dump(assembly_dict, f, indent=4)
 
 
 if __name__ == "__main__":
     profiler = cProfile.Profile()
     profiler.enable()
-    main()
+    client = Client()
+    get_random_urdf(f"{SCRIPT_DIRECTORY}{DATA_DIRECTORY}/assemblies.parquet", client)
     profiler.disable()
     stats = pstats.Stats(profiler).sort_stats("cumtime")
     stats.dump_stats("onshape.prof")

onshape_api/connect.py

@@ -33,6 +33,7 @@
 from onshape_api.models.element import Element
 from onshape_api.models.mass import MassProperties
 from onshape_api.models.variable import Variable
+from onshape_api.utilities.helpers import get_sanitized_name
 
 __all__ = ["Client", "BASE_URL", "HTTP"]
 
@@ -213,7 +214,10 @@ def get_document_metadata(self, did: str) -> DocumentMetaData:
             """
             raise ValueError(f"Access forbidden for document: {did}")
 
-        return DocumentMetaData.model_validate(res.json())
+        _document = DocumentMetaData.model_validate(res.json())
+        _document.name = get_sanitized_name(_document.name)
+
+        return _document
 
     def get_elements(self, did: str, wtype: str, wid: str) -> dict[str, Element]:
         """
@@ -327,8 +331,69 @@ def set_variables(self, did: str, wid: str, eid: str, variables: dict[str, str])
             body=payload,
         )
 
+    def get_assembly_name(
+        self,
+        did: str,
+        wtype: str,
+        wid: str,
+        eid: str,
+        configuration: str = "default",
+    ) -> str:
+        """
+        Get assembly name for a specified document / workspace / assembly.
+
+        Args:
+            did: The unique identifier of the document.
+            wtype: The type of workspace.
+            wid: The unique identifier of the workspace.
+            eid: The unique identifier of the assembly.
+
+        Returns:
+            str: Assembly name
+
+        Examples:
+            >>> assembly_name = client.get_assembly_name(
+            ...     did="a1c1addf75444f54b504f25c",
+            ...     wtype="w",
+            ...     wid="0d17b8ebb2a4c76be9fff3c7",
+            ...     eid="a86aaf34d2f4353288df8812"
+            ... )
+            >>> print(assembly_name)
+            "Assembly Name"
+        """
+        _request_path = "/api/metadata/d/" + did + "/" + wtype + "/" + wid + "/e/" + eid
+        result_json = self.request(
+            HTTP.GET,
+            _request_path,
+            query={
+                "inferMetadataOwner": "false",
+                "includeComputedProperties": "false",
+                "includeComputedAssemblyProperties": "false",
+                "thumbnail": "false",
+                "configuration": configuration,
+            },
+            log_response=False,
+        ).json()
+
+        name = None
+        try:
+            name = result_json["properties"][0]["value"]
+            name = get_sanitized_name(name)
+
+        except KeyError:
+            LOGGER.warning(f"Assembly name not found for document: {did}")
+
+        return name
+
     def get_assembly(
-        self, did: str, wtype: str, wid: str, eid: str, configuration: str = "default", log_response: bool = True
+        self,
+        did: str,
+        wtype: str,
+        wid: str,
+        eid: str,
+        configuration: str = "default",
+        log_response: bool = True,
+        with_meta_data: bool = False,
     ) -> tuple[Assembly, dict]:
         """
         Get assembly data for a specified document / workspace / assembly.
@@ -340,6 +405,7 @@ def get_assembly(
             eid: The unique identifier of the assembly.
             configuration: The configuration of the assembly.
             log_response: Log the response from the API request.
+            with_meta_data: Include meta data in the assembly data.
 
         Returns:
             Assembly: Assembly object containing the assembly data
@@ -394,6 +460,11 @@ def get_assembly(
         _document = Document(did=did, wtype=wtype, wid=wid, eid=eid)
         _assembly.document = _document
 
+        if with_meta_data:
+            _assembly.name = self.get_assembly_name(did, wtype, wid, eid, configuration)
+            _document_meta_data = self.get_document_metadata(did)
+            _assembly.document.name = _document_meta_data.name
+
         return _assembly, _assembly_json
 
     def download_stl(self, did: str, wid: str, eid: str, partID: str, buffer: BinaryIO) -> BinaryIO:

onshape_api/models/assembly.py

@@ -1176,6 +1176,7 @@ class Assembly(BaseModel):
 
     Custom Attributes:
         document (Union[Document, None]): The document object associated with the assembly. Defaults to None.
+        name (Union[str, None]): The name of the assembly. Defaults to None.
 
 
     Examples:
@@ -1219,3 +1220,4 @@ class Assembly(BaseModel):
     partStudioFeatures: list[dict] = Field(..., description="A list of part studio features in the document.")
 
     document: Union[Document, None] = Field(None, description="The document associated with the assembly.")
+    name: Union[str, None] = Field(None, description="The name of the assembly.")

onshape_api/parse.py

@@ -277,9 +277,13 @@ def get_mates_and_relations(
     Args:
         assembly: The Onshape assembly object to use for extracting mates.
         subassembly_map: Mapping of subassembly IDs to their corresponding subassembly objects.
+        id_to_name_map: Mapping of instance IDs to their corresponding sanitized names. This can be obtained
+            by calling the `get_instances` function.
 
     Returns:
-        A dictionary mapping occurrence paths to their corresponding mate data.
+        A tuple containing:
+        - A dictionary mapping mate IDs to their corresponding mate feature data.
+        - A dictionary mapping mate relation IDs to their corresponding mate relation feature data.
 
     Examples:
         >>> assembly = Assembly(...)