Added support for mate relations, joint name references still need to…

… be sorted out. Made the urdf pretty looking.

.gitignore

@@ -174,3 +174,4 @@ benchmark/**/*.urdf
 benchmark/**/*.stl
 benchmark/**/*.png
 benchmark/**/*.prof
+benchmark/**/*.json

benchmark/relations.py

@@ -0,0 +1,51 @@
+import json
+import os
+
+import onshape_api as opa
+from onshape_api.models.document import Document
+
+SCRIPT_DIRECTORY = os.path.dirname(os.path.realpath(__file__))
+
+if __name__ == "__main__":
+    client = opa.Client()
+
+    json_path = SCRIPT_DIRECTORY + "/arbor_press.json"
+    json_data = json.load(open(json_path))
+
+    assembly = opa.Assembly.model_validate(json_data)
+    document = client.get_document_metadata(assembly.rootAssembly.documentId)
+    assembly.document = Document(
+        did=assembly.rootAssembly.documentId,
+        wtype=document.defaultWorkspace.type.shorthand,
+        wid=document.defaultWorkspace.id,
+        eid=assembly.rootAssembly.elementId,
+        name=document.name,
+    )
+
+    print(assembly.document.url)
+
+    instances, id_to_name_map = opa.get_instances(assembly)
+    occurences = opa.get_occurences(assembly, id_to_name_map)
+    parts = opa.get_parts(assembly, client, instances)
+    subassemblies = opa.get_subassemblies(assembly, instances)
+    mates, relations = opa.get_mates_and_relations(assembly, subassemblies, id_to_name_map)
+
+    graph, root_node = opa.create_graph(
+        occurences=occurences,
+        instances=instances,
+        parts=parts,
+        mates=mates,
+    )
+
+    links, joints = opa.get_urdf_components(
+        assembly=assembly,
+        graph=graph,
+        root_node=root_node,
+        parts=parts,
+        mates=mates,
+        relations=relations,
+        client=client,
+    )
+
+    robot = opa.Robot(name=document.name, links=links, joints=joints)
+    robot.save("robot.urdf")

onshape_api/__init__.py

@@ -11,5 +11,9 @@ def get_version() -> str:
 __version__: str = get_version()
 
 from onshape_api.connect import *  # noqa: F403 E402
+from onshape_api.graph import *  # noqa: F403 E402
 from onshape_api.log import *  # noqa: F403 E402
+from onshape_api.mesh import *  # noqa: F403 E402
 from onshape_api.models import *  # noqa: F403 E402
+from onshape_api.parse import *  # noqa: F403 E402
+from onshape_api.urdf import *  # noqa: F403 E402

onshape_api/data/preprocess.py

@@ -3,6 +3,7 @@
 import re
 from functools import partial
 
+import numpy as np
 import pandas as pd
 from tqdm import tqdm
 
@@ -90,6 +91,33 @@ def get_assembly_df(automate_assembly_df: pd.DataFrame, client: Client, chunk_si
     return assembly_df
 
 
+def process_all_checkpoints(client: Client):
+    assemblies_df = pd.DataFrame()
+    MAX_CHECKPOINTS = 256
+
+    for i in range(MAX_CHECKPOINTS):
+        checkpoint_path = f"assemblies_checkpoint_{i}.parquet"
+        if os.path.exists(checkpoint_path):
+            assembly_df = pd.read_parquet(checkpoint_path, engine="pyarrow")
+            LOGGER.info(f"Processing checkpoint: {checkpoint_path} with {assembly_df.shape[0]} rows")
+            assembly_df.dropna(subset=["documentId", "elementId"], inplace=True)
+
+            assembly_df["elementId"] = assembly_df["elementId"].apply(
+                lambda x: ", ".join(x) if isinstance(x, (list, np.ndarray)) else x
+            )
+            # drop all duplicate entries
+            assembly_df.drop_duplicates(subset=["documentId", "elementId"], inplace=True)
+            LOGGER.info(f"Checkpoint {checkpoint_path} processed with {assembly_df.shape[0]} rows")
+            LOGGER.info("--" * 20)
+            assemblies_df = pd.concat([assemblies_df, assembly_df], ignore_index=True)
+
+    assemblies_df["elementId"] = assemblies_df["elementId"].apply(lambda x: x.split(", ") if isinstance(x, str) else x)
+
+    # now for every elementId in the list, we will have a separate row
+    assemblies_df = assemblies_df.explode("elementId")
+    assemblies_df.to_parquet("assemblies.parquet", engine="pyarrow")
+
+
 def save_all_jsons(client: Client):
     if not os.path.exists("assemblies.parquet"):
         # TODO: Re-process the automate data if assemblies.parquet is empty
@@ -104,30 +132,28 @@ def save_all_jsons(client: Client):
 
     for _, row in tqdm(assembly_df.iterrows(), total=len(assembly_df)):
         try:
-            for element_id in row["elementId"]:
-                _, assembly_json = client.get_assembly(
-                    did=row["documentId"],
-                    wtype=row["wtype"],
-                    wid=row["workspaceId"],
-                    eid=element_id,
-                    log_response=False,
-                )
+            _, assembly_json = client.get_assembly(
+                did=row["documentId"],
+                wtype=row["wtype"],
+                wid=row["workspaceId"],
+                eid=row["elementId"],
+                log_response=False,
+            )
 
-                json_file_path = os.path.join(json_dir, f"{row['documentId']}_{element_id}.json")
-                with open(json_file_path, "w") as json_file:
-                    json.dump(assembly_json, json_file, indent=4)
+            json_file_path = os.path.join(json_dir, f"{row['documentId']}_{row["elementId"]}.json")
+            with open(json_file_path, "w") as json_file:
+                json.dump(assembly_json, json_file, indent=4)
 
-                LOGGER.info(f"Assembly JSON saved to {json_file_path}")
+            LOGGER.info(f"Assembly JSON saved to {json_file_path}")
 
         except Exception as e:
             LOGGER.warning(f"Error saving assembly JSON: {os.path.abspath(json_file_path)}")
-            document_url = generate_url(row["documentId"], row["wtype"], row["workspaceId"], element_id)
+            document_url = generate_url(row["documentId"], row["wtype"], row["workspaceId"], row["elementId"])
             LOGGER.warning(f"Onshape document: {document_url}")
             LOGGER.warning(f"Assembly JSON: {assembly_json}")
             LOGGER.warning(f"Element ID: {row['elementId']}")
             LOGGER.warning(e)
-
-            break
+            continue
 
 
 def validate_assembly_json(json_file_path: str):

onshape_api/models/assembly.py

@@ -705,6 +705,9 @@ class MateGroupFeatureData(BaseModel):
         occurrences (list[MateGroupFeatureOccurrence]): A list of occurrences in the mate group feature.
         name (str): The name of the mate group feature.
 
+    Custom Attributes:
+        id (str): The unique identifier of the feature.
+
     Examples:
         >>> MateGroupFeatureData(
         ...     occurrences=[
@@ -729,6 +732,8 @@ class MateGroupFeatureData(BaseModel):
     )
     name: str = Field(..., description="The name of the mate group feature.")
 
+    id: str = Field(None, description="The unique identifier of the feature.")
+
 
 class MateConnectorFeatureData(BaseModel):
     """
@@ -755,6 +760,9 @@ class MateConnectorFeatureData(BaseModel):
         occurrence (list[str]): A list of identifiers for the occurrences involved in the mate connector.
         name (str): The name of the mate connector feature.
 
+    Custom Attributes:
+        id (str): The unique identifier of the feature.
+
     Examples:
         >>> MateConnectorFeatureData(
         ...     mateConnectorCS=MatedCS(
@@ -778,9 +786,13 @@ class MateConnectorFeatureData(BaseModel):
         )
     """
 
-    mateConnectorCS: MatedCS
-    occurrence: list[str]
-    name: str
+    mateConnectorCS: MatedCS = Field(..., description="The coordinate system used for the mate connector.")
+    occurrence: list[str] = Field(
+        ..., description="A list of identifiers for the occurrences involved in the mate connector."
+    )
+    name: str = Field(..., description="The name of the mate connector feature.")
+
+    id: str = Field(None, description="The unique identifier of the feature.")
 
 
 class MateRelationFeatureData(BaseModel):
@@ -811,9 +823,13 @@ class MateRelationFeatureData(BaseModel):
         relationType (RelationType): The type of mate relation.
         mates (list[MateRelationMate]): A list of mate relations.
         reverseDirection (bool): Indicates if the direction of the mate relation is reversed.
-        relationRatio (Union[float, None]): The ratio of the mate relation. Defaults to None.
+        relationRatio (Union[float, None]): The ratio of the GEAR mate relation. Defaults to None.
+        relationLength (Union[float, None]): The length of the RACK_AND_PINION mate relation. Defaults to None.
         name (str): The name of the mate relation feature.
 
+    Custom Attributes:
+        id (str): The unique identifier of the feature.
+
     Examples:
         >>> MateRelationFeatureData(
         ...     relationType=RelationType.GEAR,
@@ -834,9 +850,16 @@ class MateRelationFeatureData(BaseModel):
     relationType: RelationType = Field(..., description="The type of mate relation.")
     mates: list[MateRelationMate] = Field(..., description="A list of mate relations.")
     reverseDirection: bool = Field(..., description="Indicates if the direction of the mate relation is reversed.")
-    relationRatio: Union[float, None] = Field(None, description="The ratio of the mate relation. Defaults to None.")
+    relationRatio: Union[float, None] = Field(
+        None, description="The ratio of the GEAR mate relation. Defaults to None."
+    )
+    relationLength: Union[float, None] = Field(
+        None, description="The length of the RACK_AND_PINION mate relation. Defaults to None."
+    )
     name: str = Field(..., description="The name of the mate relation feature.")
 
+    id: str = Field(None, description="The unique identifier of the feature.")
+
 
 class MateFeatureData(BaseModel):
     """
@@ -877,6 +900,9 @@ class MateFeatureData(BaseModel):
         mateType (MateType): The type of mate.
         name (str): The name of the mate feature.
 
+    Custom Attributes:
+        id (str): The unique identifier of the feature.
+
     Examples:
         >>> MateFeatureData(
         ...     matedEntities=[...],
@@ -895,6 +921,8 @@ class MateFeatureData(BaseModel):
     mateType: MateType = Field(..., description="The type of mate.")
     name: str = Field(..., description="The name of the mate feature.")
 
+    id: str = Field(None, description="The unique identifier of the feature.")
+
 
 class AssemblyFeature(BaseModel):
     """

onshape_api/models/document.py

@@ -155,6 +155,7 @@ class Document(BaseModel):
         wtype: The type of workspace (w, v, m)
         wid: The unique identifier of the workspace
         eid: The unique identifier of the element
+        name: The name of the document
 
     Methods:
         from_url: Create a Document instance from an Onshape URL
@@ -181,6 +182,7 @@ class Document(BaseModel):
     wtype: str = Field(..., description="The type of workspace (w, v, m)")
     wid: str = Field(..., description="The unique identifier of the workspace")
     eid: str = Field(..., description="The unique identifier of the element")
+    name: str = Field(None, description="The name of the document")
 
     def __init__(self, **data):
         super().__init__(**data)

onshape_api/models/robot.py

@@ -11,6 +11,8 @@
 from dataclasses import dataclass
 from pathlib import Path
 
+from defusedxml import minidom
+
 from onshape_api.models.joint import BaseJoint
 from onshape_api.models.link import Link
 
@@ -77,4 +79,8 @@ def save(self, path: str | Path | io.StringIO) -> None:
 
         tree = ET.ElementTree(self.to_xml())
         if isinstance(path, (str, Path)):
-            tree.write(path, encoding="unicode", xml_declaration=True)
+            xml_str = ET.tostring(tree.getroot(), encoding="unicode")
+            pretty_xml_str = minidom.parseString(xml_str).toprettyxml(indent="    ")
+            print(pretty_xml_str)
+            with open(path, "w", encoding="utf-8") as f:
+                f.write(pretty_xml_str)

onshape_api/parse.py

@@ -11,11 +11,11 @@
 from onshape_api.log import LOGGER
 from onshape_api.models.assembly import (
     Assembly,
-    AssemblyFeature,
     AssemblyFeatureType,
     AssemblyInstance,
     InstanceType,
     MateFeatureData,
+    MateRelationFeatureData,
     Occurrence,
     Part,
     PartInstance,
@@ -30,8 +30,8 @@
 SUBASSEMBLY_JOINER = "-SUB-"
 MATE_JOINER = "-MATE-"
 
-CHILD = 0
-PARENT = 1
+CHILD = 1
+PARENT = 0
 
 
 def get_instances(assembly: Assembly) -> tuple[dict[str, Union[PartInstance, AssemblyInstance]], dict[str, str]]:
@@ -265,13 +265,13 @@ def join_mate_occurences(parent: list[str], child: list[str], prefix: Optional[s
         return f"{parent_occurence}{MATE_JOINER}{child_occurence}"
 
 
-def get_mates(
+def get_mates_and_relations(
     assembly: Assembly,
     subassembly_map: dict[str, SubAssembly],
     id_to_name_map: dict[str, str],
-) -> dict[str, AssemblyFeature]:
+) -> tuple[dict[str, MateFeatureData], dict[str, MateRelationFeatureData]]:
     """
-    Get mates of the assembly.
+    Get mates and relations of an Onshape assembly.
 
     Args:
         assembly: The Onshape assembly object to use for extracting mates.
@@ -282,20 +282,29 @@ def get_mates(
 
     Examples:
         >>> assembly = Assembly(...)
-        >>> get_mates(assembly)
-        {
+        >>> get_mates_and_relations(assembly)
+        ({
             "subassembly1-SUB-part1-MATE-subassembly2-SUB-part2": AssemblyFeature(...),
-            "part1-MATE-part2": AssemblyFeature(...),
-        }
+            "part1-MATE-part2": MateFeatureData(...),
+        },
+        {
+            "MuwOg31fsdH/5O2nX": MateRelationFeatureData(...),
+        })
     """
 
     def traverse_assembly(
         root: Union[RootAssembly, SubAssembly], subassembly_prefix: Optional[str] = None
-    ) -> dict[str, MateFeatureData]:
+    ) -> tuple[dict[str, MateFeatureData], dict[str, MateRelationFeatureData]]:
         _mates_map = {}
+        _relations_map = {}
 
         for feature in root.features:
-            if feature.featureType == AssemblyFeatureType.MATE and not feature.suppressed:
+            feature.featureData.id = feature.id
+
+            if feature.suppressed:
+                continue
+
+            if feature.featureType == AssemblyFeatureType.MATE:
                 if len(feature.featureData.matedEntities) < 2:
                     # TODO: will there be features with just one mated entity?
                     LOGGER.warning(f"Invalid mate feature: {feature}")
@@ -319,11 +328,17 @@ def traverse_assembly(
                     )
                 ] = feature.featureData
 
-        return _mates_map
+            elif feature.featureType == AssemblyFeatureType.MATERELATION:
+                child_joint_id = feature.featureData.mates[CHILD].featureId
+                _relations_map[child_joint_id] = feature.featureData
+
+        return _mates_map, _relations_map
 
-    mates_map = traverse_assembly(assembly.rootAssembly)
+    mates_map, relations_map = traverse_assembly(assembly.rootAssembly)
 
     for key, subassembly in subassembly_map.items():
-        mates_map.update(traverse_assembly(subassembly, key))
+        sub_mates_map, sub_relations_map = traverse_assembly(subassembly, key)
+        mates_map.update(sub_mates_map)
+        relations_map.update(sub_relations_map)
 
-    return mates_map
+    return mates_map, relations_map