Major fixes to subassembly parsing. Instance map is utilized now and …

…fixed a dumb mistake in mates parsing.
neurobionics · Nov 20, 2024 · 09e0b1c · 09e0b1c
1 parent 7abdecb
commit 09e0b1c
Show file tree

Hide file tree

Showing 12 changed files with 100 additions and 67 deletions.
diff --git a/benchmark/assembly.py b/benchmark/assembly.py
@@ -6,7 +6,7 @@
 import pandas as pd
 
 from onshape_api.connect import Client
-from onshape_api.graph import create_graph, save_graph
+from onshape_api.graph import create_graph, plot_graph
 from onshape_api.models.robot import Robot
 from onshape_api.parse import (
     get_instances,
@@ -54,7 +54,7 @@ def get_random_urdf(data_path: str, client: Client):
         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"{assembly_robot_name}.png")
+        plot_graph(graph, f"{assembly_robot_name}.png")
 
         links, joints = get_urdf_components(
             assembly=assembly,

diff --git a/benchmark/connectors.py b/benchmark/connectors.py
@@ -39,7 +39,7 @@
         directed=True,
         use_user_defined_root=True,
     )
-    opa.save_graph(graph, f"{assembly_robot_name}.png")
+    opa.plot_graph(graph, f"{assembly_robot_name}.png")
 
     links, joints = opa.get_urdf_components(
         assembly=assembly,

diff --git a/benchmark/relations.py b/benchmark/relations.py
@@ -38,7 +38,7 @@
         parts=parts,
         mates=mates,
     )
-    opa.save_graph(graph, f"{assembly_robot_name}.png")
+    opa.plot_graph(graph, f"{assembly_robot_name}.png")
 
     links, joints = opa.get_urdf_components(
         assembly=assembly,

diff --git a/benchmark/robots.py b/benchmark/robots.py
@@ -8,9 +8,13 @@
 if __name__ == "__main__":
     client = opa.Client()
     # robot = https://cad.onshape.com/documents/a8f62e825e766a6512320ceb/w/b9099bcbdc92e6d6c810f0b7/e/f5b0475edd5ad0193d280fc4
+    # robot dog = https://cad.onshape.com/documents/d0223bce364d259e80667122/w/b52c33333c8553dce379aac6/e/57728d0a8bc87b7b065e43be
+    # simple robot dog = https://cad.onshape.com/documents/64d7b47821f3f5c91e3cd128/w/051d83c286bca38e8952dd84/e/ba886678bddf9de9c01723c8
+
+    # test-nested-subassemblies = https://cad.onshape.com/documents/8c7a1c45e27a40a5b6e44d92/w/9c50078d1ac7106985359fe8/e/8c0e0762c95eb6e8b2f4b1f1
 
     document = Document.from_url(
-        "https://cad.onshape.com/documents/cf6b852d2c88d661ac2e17e8/w/c842455c29cc878dc48bdc68/e/b5e293d409dd0b88596181ef"
+        "https://cad.onshape.com/documents/8c7a1c45e27a40a5b6e44d92/w/9c50078d1ac7106985359fe8/e/8c0e0762c95eb6e8b2f4b1f1"
     )
     assembly, _ = client.get_assembly(
         did=document.did,
@@ -22,21 +26,23 @@
 
     opa.LOGGER.info(assembly.document.url)
     assembly_robot_name = f"{assembly.document.name + '-' + assembly.name}"
+    opa.save_model_as_json(assembly, f"{assembly_robot_name}.json")
 
     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)
 
+    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,
         use_user_defined_root=False,
     )
-    opa.save_graph(graph, f"{assembly_robot_name}.png")
+    opa.plot_graph(graph)
 
     links, joints = opa.get_urdf_components(
         assembly=assembly,

diff --git a/docs/tutorials/edit.md b/docs/tutorials/edit.md
@@ -97,11 +97,11 @@ mates, relations = get_mates_and_relations(assembly, subassembly_map=subassembli
 Generate a graph visualization of the assembly structure:
 
 ```python
-from onshape_api.graph import create_graph, save_graph
+from onshape_api.graph import create_graph, plot_graph
 
 # Create and save the assembly graph
 graph, root_node = create_graph(occurences=occurences, instances=instances, parts=parts, mates=mates)
-save_graph(graph, "bike.png")
+plot_graph(graph, "bike.png")
 ```
 
 <img src="bike-graph.png" alt="Bike Graph" style="width: 100%;">

diff --git a/docs/tutorials/export.md b/docs/tutorials/export.md
@@ -86,7 +86,7 @@ graph, root_node = osa.create_graph(
     parts=parts,
     mates=mates,
 )
-osa.save_graph(graph, f"{assembly.document.name + '-' + assembly.name}.png")
+osa.plot_graph(graph, f"{assembly.document.name + '-' + assembly.name}.png")
 ```
 
 This will save a PNG file of the assembly graph in your current working directory.

diff --git a/examples/bike/main.py b/examples/bike/main.py
@@ -1,5 +1,5 @@
 import onshape_api as osa
-from onshape_api.graph import create_graph, save_graph
+from onshape_api.graph import create_graph, plot_graph
 from onshape_api.models.robot import Robot
 from onshape_api.parse import (
     get_instances,
@@ -33,7 +33,7 @@
 mates, relations = get_mates_and_relations(assembly, subassembly_map=subassemblies, id_to_name_map=id_to_name_map)
 
 graph, root_node = create_graph(occurences=occurences, instances=instances, parts=parts, mates=mates)
-save_graph(graph, "bike.png")
+plot_graph(graph, "bike.png")
 
 links, joints = get_urdf_components(assembly, graph, root_node, parts, mates, relations, client)
 robot = Robot(name="bike", links=links, joints=joints)

diff --git a/onshape_api/__init__.py b/onshape_api/__init__.py
@@ -17,3 +17,4 @@ def get_version() -> str:
 from onshape_api.models import *  # noqa: F403 E402
 from onshape_api.parse import *  # noqa: F403 E402
 from onshape_api.urdf import *  # noqa: F403 E402
+from onshape_api.utilities import *  # noqa: F403 E402
diff --git a/onshape_api/connect.py b/onshape_api/connect.py
@@ -507,6 +507,7 @@ def download_stl(
             >>> raw_mesh = stl.mesh.Mesh.from_file(None, fh=buffer)
             >>> raw_mesh.save("mesh.stl")
         """
+        # TODO: version id seems to always work, should this default behavior be changed?
 
         req_headers = {"Accept": "application/vnd.onshape.v1+octet-stream"}
         _request_path = (
@@ -591,17 +592,17 @@ def get_mass_property(
                 principalAxes=[...]
             )
         """
+        # TODO: version id seems to always work, should this default behavior be changed?
         _request_path = (
             f"/api/parts/d/{did}/{wtype}/"
             f"{wid if wtype == WorkspaceType.W else vid}/e/{eid}/partid/{partID}/massproperties"
         )
-        res = self.request(HTTP.GET, _request_path, {"useMassPropertiesOverrides": True})
+        res = self.request(HTTP.GET, _request_path, {"useMassPropertiesOverrides": True}, log_response=False)
 
         if res.status_code == 404:
             # TODO: There doesn't seem to be a way to assign material to a part currently
             # It is possible that the workspace got deleted
             if vid and wtype == WorkspaceType.W:
-                print("Trying to get mass properties from a version workspace")
                 return self.get_mass_property(did, wid, eid, partID, vid, WorkspaceType.V.value)
 
             raise ValueError(f"Part: {

diff --git a/onshape_api/graph.py b/onshape_api/graph.py
@@ -4,7 +4,7 @@
 """
 
 import random
-from typing import Union
+from typing import Optional, Union
 
 import matplotlib.pyplot as plt
 import networkx as nx
@@ -21,48 +21,44 @@
 from onshape_api.parse import MATE_JOINER
 
 
-def show_graph(graph: nx.Graph) -> None:
+def plot_graph(graph: Union[nx.Graph, nx.DiGraph], file_name: Optional[str] = None) -> None:
     """
-    Display the graph using networkx and matplotlib.
+    Display the graph using networkx and matplotlib, or save it as an image file.
 
     Args:
-        graph: The graph to display.
+        graph: The graph to display or save.
+        file_name: The name of the image file to save. If None, the graph will be displayed.
 
     Examples:
         >>> graph = nx.Graph()
-        >>> show_graph(graph)
+        >>> plot_graph(graph)
+        >>> plot_graph(graph, "graph.png")
     """
-    nx.draw_circular(graph, with_labels=True)
-    plt.show()
-
-
-def save_graph(graph: Union[nx.Graph, nx.DiGraph], file_name: str) -> None:
-    """
-    Save the graph as an image file.
-
-    Args:
-        graph: The graph to save.
-        file_name: The name of the image file.
-
-    Examples:
-        >>> graph = nx.Graph()
-        >>> save_graph(graph, "graph.png")
-    """
-
     colors = [f"#{random.randint(0, 0xFFFFFF):06x}" for _ in range(len(graph.nodes))]  # noqa: S311
     plt.figure(figsize=(8, 8))
-    pos = nx.circular_layout(graph)
-    nx.draw(
-        graph,
-        pos,
-        with_labels=True,
-        arrows=True,
-        node_color=colors,
-        edge_color="white",
-        font_color="white",
-    )
-    plt.savefig(file_name, transparent=True)
-    plt.close()
+    pos = nx.spring_layout(graph)
+
+    if file_name:
+        nx.draw(
+            graph,
+            pos,
+            with_labels=True,
+            arrows=True,
+            node_color=colors,
+            edge_color="white",
+            font_color="white",
+        )
+        plt.savefig(file_name, transparent=True)
+        plt.close()
+    else:
+        nx.draw(
+            graph,
+            pos,
+            with_labels=True,
+            arrows=True,
+            node_color=colors,
+        )
+        plt.show()
 
 
 def get_root_node(graph: nx.DiGraph) -> str:

diff --git a/onshape_api/parse.py b/onshape_api/parse.py
@@ -182,14 +182,17 @@ def get_subassemblies(
             "subassembly2": SubAssembly(...),
         }
     """
-    subassembly_map = {}
+    subassembly_map: dict[str, SubAssembly] = {}
+    subassembly_instance_map: dict[str, list[str]] = {}
+
+    for key, instance in instance_map.items():
+        if instance.type == InstanceType.ASSEMBLY:
+            subassembly_instance_map.setdefault(instance.uid, []).append(key)
 
-    subassembly_instance_map = {
-        instance.uid: key for key, instance in instance_map.items() if instance.type == InstanceType.ASSEMBLY
-    }
     for subassembly in assembly.subAssemblies:
         if subassembly.uid in subassembly_instance_map:
-            subassembly_map[subassembly_instance_map[subassembly.uid]] = subassembly
+            for key in subassembly_instance_map[subassembly.uid]:
+                subassembly_map[key] = subassembly
 
     return subassembly_map
 
@@ -275,7 +278,7 @@ def join_mate_occurences(parent: list[str], child: list[str], prefix: Optional[s
         return f"{parent_occurence}{MATE_JOINER}{child_occurence}"
 
 
-def get_mates_and_relations(
+def get_mates_and_relations(  # noqa: C901
     assembly: Assembly,
     subassembly_map: dict[str, SubAssembly],
     id_to_name_map: dict[str, str],
@@ -324,15 +327,18 @@ def traverse_assembly(
                     LOGGER.warning(f"Invalid mate feature: {feature}")
                     continue
 
-                child_occurences = [
-                    id_to_name_map[path] for path in feature.featureData.matedEntities[CHILD].matedOccurrence
-                ]
-                parent_occurences = [
-                    id_to_name_map[path] for path in feature.featureData.matedEntities[PARENT].matedOccurrence
-                ]
-
-                feature.featureData.matedEntities[CHILD].matedOccurrence = child_occurences
-                feature.featureData.matedEntities[PARENT].matedOccurrence = parent_occurences
+                try:
+                    child_occurences = [
+                        id_to_name_map[path] for path in feature.featureData.matedEntities[CHILD].matedOccurrence
+                    ]
+                    parent_occurences = [
+                        id_to_name_map[path] for path in feature.featureData.matedEntities[PARENT].matedOccurrence
+                    ]
+                except KeyError as e:
+                    LOGGER.warning(e)
+                    LOGGER.warning(f"Key not found in {id_to_name_map.keys()}")
+                    LOGGER.warning(f"Occurrence path not found for mate feature: {feature}")
+                    continue
 
                 _mates_map[
                     join_mate_occurences(
@@ -344,17 +350,14 @@ def traverse_assembly(
 
             elif feature.featureType == AssemblyFeatureType.MATERELATION:
                 if feature.featureData.relationType == RelationType.SCREW:
-                    # Screw relations only have one reference entity
                     child_joint_id = feature.featureData.mates[0].featureId
                 else:
-                    # TODO: Verify mate relation convention
                     child_joint_id = feature.featureData.mates[RELATION_CHILD].featureId
 
                 _relations_map[child_joint_id] = feature.featureData
 
             elif feature.featureType == AssemblyFeatureType.MATECONNECTOR:
-                # Mate connectors' MatedCS data is already included in the MateFeatureData
-                # TODO: This might not be true for all cases?
+                # TODO: Mate connectors' MatedCS data is already included in the MateFeatureData
                 pass
 
         return _mates_map, _relations_map

diff --git a/onshape_api/utilities/helpers.py b/onshape_api/utilities/helpers.py
@@ -11,13 +11,39 @@
 """
 
 import hashlib
+import json
 import os
 import random
 from xml.sax.saxutils import escape
 
+from pydantic import BaseModel
+
 from onshape_api.log import LOGGER
 
 
+def save_model_as_json(model: BaseModel, file_path: str) -> None:
+    """
+    Save a Pydantic model as a JSON file
+
+    Args:
+        model (BaseModel): Pydantic model to save
+        file_path (str): File path to save JSON file
+
+    Returns:
+        None
+
+    Examples:
+        >>> class TestModel(BaseModel):
+        ...     a: int
+        ...     b: str
+        ...
+        >>> save_model_as_json(TestModel(a=1, b="hello"), "test.json")
+    """
+
+    with open(file_path, "w") as file:
+        json.dump(model.model_dump(), file, indent=4)
+
+
 def xml_escape(unescaped: str) -> str:
     """
     Escape XML characters in a string