From 9b44a402f1e14b063a141290441ffaf9800d3da4 Mon Sep 17 00:00:00 2001 From: imsenthur Date: Sun, 15 Dec 2024 20:46:44 -0500 Subject: [PATCH] Restructured convenience methods and created from_url and from_urdf class methods within Robot class. --- examples/edit/main.py | 8 +- examples/export/main.py | 56 ++------ onshape_api/connect.py | 21 ++- onshape_api/parse.py | 59 ++++----- onshape_api/robot.py | 212 +++++++++++++++++++++++++++++-- onshape_api/urdf.py | 104 +-------------- onshape_api/utilities/helpers.py | 16 ++- 7 files changed, 281 insertions(+), 195 deletions(-) diff --git a/examples/edit/main.py b/examples/edit/main.py index 6c61cbc..903c858 100644 --- a/examples/edit/main.py +++ b/examples/edit/main.py @@ -1,5 +1,5 @@ from onshape_api.connect import Client -from onshape_api.graph import create_graph, plot_graph +from onshape_api.graph import create_graph from onshape_api.log import LOGGER, LogLevel from onshape_api.models.document import Document from onshape_api.parse import ( @@ -8,7 +8,7 @@ get_parts, get_subassemblies, ) -from onshape_api.urdf import get_robot +from onshape_api.robot import get_robot if __name__ == "__main__": LOGGER.set_file_name("edit.log") @@ -38,6 +38,6 @@ graph, root_node = create_graph(occurrences=occurrences, instances=instances, parts=parts, mates=mates) robot = get_robot(assembly, graph, root_node, parts, mates, relations, client, "test") - robot.show() - plot_graph(graph, "bike.png") + robot.show_tree() + robot.show_graph("bike.png") robot.save() diff --git a/examples/export/main.py b/examples/export/main.py index e0ec9da..b7bffbc 100644 --- a/examples/export/main.py +++ b/examples/export/main.py @@ -1,57 +1,23 @@ -import os - from onshape_api.connect import Client -from onshape_api.graph import create_graph, plot_graph from onshape_api.log import LOGGER, LogLevel -from onshape_api.models.document import Document -from onshape_api.parse import get_instances, get_mates_and_relations, get_parts, get_subassemblies -from onshape_api.urdf import get_robot +from onshape_api.robot import Robot from onshape_api.utilities.helpers import save_model_as_json -SCRIPT_DIRECTORY = os.path.dirname(os.path.realpath(__file__)) - if __name__ == "__main__": - LOGGER.set_file_name("export.log") + LOGGER.set_file_name("ballbot.log") LOGGER.set_stream_level(LogLevel.INFO) client = Client() - document = Document.from_url( - "https://cad.onshape.com/documents/1f42f849180e6e5c9abfce52/w/0c00b6520fac5fada24b2104/e/c96b40ef586e60c182f41d29" - ) - assembly = client.get_assembly( - did=document.did, - wtype=document.wtype, - wid=document.wid, - eid=document.eid, - ) - - LOGGER.info(assembly.document.url) - assembly_robot_name = f"{assembly.document.name + '-' + assembly.name}" - save_model_as_json(assembly, f"{assembly_robot_name}.json") - - instances, occurrences, id_to_name_map = get_instances(assembly) - subassemblies, rigid_subassemblies = get_subassemblies(assembly, client, instances) - - parts = get_parts(assembly, rigid_subassemblies, client, instances) - mates, relations = get_mates_and_relations(assembly, subassemblies, rigid_subassemblies, id_to_name_map, parts) - - graph, root_node = create_graph( - occurrences=occurrences, - instances=instances, - parts=parts, - mates=mates, + robot = Robot.from_url( + name="ballbot", + url="https://cad.onshape.com/documents/1f42f849180e6e5c9abfce52/w/0c00b6520fac5fada24b2104/e/c96b40ef586e60c182f41d29", + client=client, + max_depth=0, use_user_defined_root=False, ) - plot_graph(graph, f"{assembly_robot_name}.png") - robot = get_robot( - assembly=assembly, - graph=graph, - root_node=root_node, - parts=parts, - mates=mates, - relations=relations, - client=client, - robot_name=assembly_robot_name, - ) + print(robot.assembly.model_dump(), type(robot.assembly)) + save_model_as_json(robot.assembly, "ballbot.json") + + robot.show_graph() robot.save() diff --git a/onshape_api/connect.py b/onshape_api/connect.py index 9a9ca2c..ac2dd2b 100644 --- a/onshape_api/connect.py +++ b/onshape_api/connect.py @@ -128,8 +128,7 @@ class Client: Args: env (str, default='./.env'): Path to the environment file containing the access and secret keys - log_file (str, default='./onshape_api'): Path to save the log file - log_level (int, default=1): Log level (0-4) for the logger (0=DEBUG, 1=INFO, 2=WARNING, 3=ERROR, 4=CRITICAL) + base_url (str, default='https://cad.onshape.com'): Base URL for the Onshape API Methods: get_document_metadata: Get details for a specified document. @@ -144,8 +143,6 @@ class Client: Examples: >>> client = Client( ... env=".env", - ... log_file="./onshape_api", - ... log_level=1 ... ) >>> document_meta_data = client.get_document_metadata("document_id") """ @@ -171,6 +168,18 @@ def __init__(self, env: str = "./.env", base_url: str = BASE_URL): self._access_key, self._secret_key = load_env_variables(env) LOGGER.info(f"Onshape API initialized with env file: {env}") + def set_base_url(self, base_url: str): + """ + Set the base URL for the Onshape API. + + Args: + base_url: Base URL for the Onshape API + + Examples: + >>> client.set_base_url("https://cad.onshape.com") + """ + self._url = base_url + def get_document_metadata(self, did: str) -> DocumentMetaData: """ Get meta data for a specified document. @@ -987,6 +996,10 @@ def _make_headers(self, method, path, query=None, headers=None): return req_headers + @property + def base_url(self): + return self._url + class Asset: """ diff --git a/onshape_api/parse.py b/onshape_api/parse.py index 5204bff..9fa7d9f 100644 --- a/onshape_api/parse.py +++ b/onshape_api/parse.py @@ -14,6 +14,7 @@ from onshape_api.log import LOGGER from onshape_api.models.assembly import ( Assembly, + AssemblyFeature, AssemblyFeatureType, AssemblyInstance, InstanceType, @@ -268,20 +269,20 @@ async def get_subassemblies_async( def get_subassemblies( assembly: Assembly, client: Client, - instance_map: dict[str, Union[PartInstance, AssemblyInstance]], + instances: dict[str, Union[PartInstance, AssemblyInstance]], ) -> tuple[dict[str, SubAssembly], dict[str, RootAssembly]]: """ Synchronous wrapper for `get_subassemblies_async`. """ - return asyncio.run(get_subassemblies_async(assembly, client, instance_map)) + return asyncio.run(get_subassemblies_async(assembly, client, instances)) async def _fetch_mass_properties_async( part: Part, key: str, client: Client, - rigid_subassembly_map: dict[str, RootAssembly], - part_map: dict[str, Part], + rigid_subassemblies: dict[str, RootAssembly], + parts: dict[str, Part], ): """ Asynchronously fetch mass properties for a part. @@ -293,9 +294,9 @@ async def _fetch_mass_properties_async( rigid_subassembly_map: Mapping of instance IDs to rigid subassemblies. part_map: The dictionary to store fetched parts. """ - if key.split(SUBASSEMBLY_JOINER)[0] not in rigid_subassembly_map: + if key.split(SUBASSEMBLY_JOINER)[0] not in rigid_subassemblies: try: - LOGGER.info(f"Fetching mass properties for part: {part.documentVersion}, {part.partId}") + LOGGER.info(f"Fetching mass properties for part: {part.uid}, {part.partId}") part.MassProperty = await asyncio.to_thread( client.get_mass_property, did=part.documentId, @@ -307,14 +308,14 @@ async def _fetch_mass_properties_async( except Exception as e: LOGGER.error(f"Failed to fetch mass properties for part {part.partId}: {e}") - part_map[key] = part + parts[key] = part async def _get_parts_async( assembly: Assembly, - rigid_subassembly_map: dict[str, RootAssembly], + rigid_subassemblies: dict[str, RootAssembly], client: Client, - instance_map: dict[str, Union[PartInstance, AssemblyInstance]], + instances: dict[str, Union[PartInstance, AssemblyInstance]], ) -> dict[str, Part]: """ Asynchronously get parts of an Onshape assembly. @@ -331,7 +332,7 @@ async def _get_parts_async( part_instance_map: dict[str, list[str]] = {} part_map: dict[str, Part] = {} - for key, instance in instance_map.items(): + for key, instance in instances.items(): if instance.type == InstanceType.PART: part_instance_map.setdefault(instance.uid, []).append(key) @@ -339,7 +340,7 @@ async def _get_parts_async( for part in assembly.parts: if part.uid in part_instance_map: for key in part_instance_map[part.uid]: - tasks.append(_fetch_mass_properties_async(part, key, client, rigid_subassembly_map, part_map)) + tasks.append(_fetch_mass_properties_async(part, key, client, rigid_subassemblies, part_map)) await asyncio.gather(*tasks) @@ -348,9 +349,9 @@ async def _get_parts_async( def get_parts( assembly: Assembly, - rigid_subassembly_map: dict[str, RootAssembly], + rigid_subassemblies: dict[str, RootAssembly], client: Client, - instance_map: dict[str, Union[PartInstance, AssemblyInstance]], + instances: dict[str, Union[PartInstance, AssemblyInstance]], ) -> dict[str, Part]: """ Get parts of an Onshape assembly. @@ -364,7 +365,7 @@ def get_parts( Returns: A dictionary mapping part IDs to their corresponding part objects. """ - return asyncio.run(_get_parts_async(assembly, rigid_subassembly_map, client, instance_map)) + return asyncio.run(_get_parts_async(assembly, rigid_subassemblies, client, instances)) def get_occurrence_name(occurrences: list[str], subassembly_prefix: Optional[str] = None) -> str: @@ -414,13 +415,13 @@ def join_mate_occurrences(parent: list[str], child: list[str], prefix: Optional[ async def build_rigid_subassembly_occurrence_map( - rigid_subassembly_map: dict[str, RootAssembly], id_to_name_map: dict[str, str], parts: dict[str, Part] + rigid_subassemblies: dict[str, RootAssembly], id_to_name_map: dict[str, str], parts: dict[str, Part] ) -> dict[str, dict[str, Occurrence]]: """ Asynchronously build a map of rigid subassembly occurrences. """ occurrence_map: dict[str, dict[str, Occurrence]] = {} - for assembly_key, rigid_subassembly in rigid_subassembly_map.items(): + for assembly_key, rigid_subassembly in rigid_subassemblies.items(): sub_occurrences: dict[str, Occurrence] = {} for occurrence in rigid_subassembly.occurrences: try: @@ -451,11 +452,11 @@ async def build_rigid_subassembly_occurrence_map( async def process_features_async( # noqa: C901 - features: list, + features: list[AssemblyFeature], parts: dict[str, Part], id_to_name_map: dict[str, str], rigid_subassembly_occurrence_map: dict[str, dict[str, Occurrence]], - rigid_subassembly_map: dict[str, RootAssembly], + rigid_subassemblies: dict[str, RootAssembly], subassembly_prefix: Optional[str], ) -> tuple[dict[str, MateFeatureData], dict[str, MateRelationFeatureData]]: """ @@ -488,7 +489,7 @@ async def process_features_async( # noqa: C901 continue # Handle rigid subassemblies - if parent_occurrences[0] in rigid_subassembly_map: + if parent_occurrences[0] in rigid_subassemblies: _occurrence = rigid_subassembly_occurrence_map[parent_occurrences[0]].get(parent_occurrences[1]) if _occurrence: parent_parentCS = MatedCS.from_tf(np.matrix(_occurrence.transform).reshape(4, 4)) @@ -496,7 +497,7 @@ async def process_features_async( # noqa: C901 feature.featureData.matedEntities[PARENT].parentCS = parent_parentCS parent_occurrences = [parent_occurrences[0]] - if child_occurrences[0] in rigid_subassembly_map: + if child_occurrences[0] in rigid_subassemblies: _occurrence = rigid_subassembly_occurrence_map[child_occurrences[0]].get(child_occurrences[1]) if _occurrence: child_parentCS = MatedCS.from_tf(np.matrix(_occurrence.transform).reshape(4, 4)) @@ -525,8 +526,8 @@ async def process_features_async( # noqa: C901 async def get_mates_and_relations_async( assembly: Assembly, - subassembly_map: dict[str, SubAssembly], - rigid_subassembly_map: dict[str, RootAssembly], + subassemblies: dict[str, SubAssembly], + rigid_subassemblies: dict[str, RootAssembly], id_to_name_map: dict[str, str], parts: dict[str, Part], ) -> tuple[dict[str, MateFeatureData], dict[str, MateRelationFeatureData]]: @@ -534,7 +535,7 @@ async def get_mates_and_relations_async( Asynchronously get mates and relations. """ rigid_subassembly_occurrence_map = await build_rigid_subassembly_occurrence_map( - rigid_subassembly_map, id_to_name_map, parts + rigid_subassemblies, id_to_name_map, parts ) mates_map, relations_map = await process_features_async( @@ -542,13 +543,13 @@ async def get_mates_and_relations_async( parts, id_to_name_map, rigid_subassembly_occurrence_map, - rigid_subassembly_map, + rigid_subassemblies, None, ) - for key, subassembly in subassembly_map.items(): + for key, subassembly in subassemblies.items(): sub_mates, sub_relations = await process_features_async( - subassembly.features, parts, id_to_name_map, rigid_subassembly_occurrence_map, rigid_subassembly_map, key + subassembly.features, parts, id_to_name_map, rigid_subassembly_occurrence_map, rigid_subassemblies, key ) mates_map.update(sub_mates) relations_map.update(sub_relations) @@ -558,8 +559,8 @@ async def get_mates_and_relations_async( def get_mates_and_relations( assembly: Assembly, - subassembly_map: dict[str, SubAssembly], - rigid_subassembly_map: dict[str, RootAssembly], + subassemblies: dict[str, SubAssembly], + rigid_subassemblies: dict[str, RootAssembly], id_to_name_map: dict[str, str], parts: dict[str, Part], ) -> tuple[dict[str, MateFeatureData], dict[str, MateRelationFeatureData]]: @@ -567,5 +568,5 @@ def get_mates_and_relations( Synchronous wrapper for `get_mates_and_relations_async`. """ return asyncio.run( - get_mates_and_relations_async(assembly, subassembly_map, rigid_subassembly_map, id_to_name_map, parts) + get_mates_and_relations_async(assembly, subassemblies, rigid_subassemblies, id_to_name_map, parts) ) diff --git a/onshape_api/robot.py b/onshape_api/robot.py index 289e04e..efca4ba 100644 --- a/onshape_api/robot.py +++ b/onshape_api/robot.py @@ -13,18 +13,39 @@ import networkx as nx from lxml import etree as ET -from onshape_api.connect import Asset +from onshape_api.connect import Asset, Client +from onshape_api.graph import create_graph, plot_graph from onshape_api.log import LOGGER +from onshape_api.models.assembly import ( + Assembly, + MateFeatureData, + MateRelationFeatureData, + Part, + RelationType, + RootAssembly, + SubAssembly, +) +from onshape_api.models.document import Document from onshape_api.models.joint import ( BaseJoint, ContinuousJoint, FixedJoint, FloatingJoint, + JointMimic, JointType, PrismaticJoint, RevoluteJoint, ) from onshape_api.models.link import Link +from onshape_api.parse import ( + MATE_JOINER, + RELATION_PARENT, + get_instances, + get_mates_and_relations, + get_parts, + get_subassemblies, +) +from onshape_api.urdf import get_joint_name, get_robot_joint, get_robot_link, get_topological_mates class RobotType(str, Enum): @@ -88,10 +109,20 @@ class Robot: """ def __init__(self, name: str, assets: Optional[dict[str, Asset]] = None, robot_type: RobotType = RobotType.URDF): - self.name = name - self.graph = nx.DiGraph() # Graph to hold links and joints - self.assets = assets - self.type = robot_type + self.name: str = name + self.graph: nx.DiGraph = nx.DiGraph() # Graph to hold links and joints + self.assets: Optional[dict[str, Asset]] = assets + self.type: RobotType = robot_type + + # Onshape assembly attributes + self.parts: dict[str, Part] = {} + self.mates: dict[str, MateFeatureData] = {} + self.relations: dict[str, MateRelationFeatureData] = {} + + self.subassemblies: dict[str, SubAssembly] = {} + self.rigid_subassemblies: dict[str, RootAssembly] = {} + + self.assembly: Optional[Assembly] = None def add_link(self, link: Link) -> None: """Add a link to the graph.""" @@ -138,7 +169,7 @@ def save(self, file_path: Optional[str] = None, download_assets: bool = True) -> LOGGER.info(f"Robot model saved to {file_path}") - def show(self) -> None: + def show_tree(self) -> None: """Display the robot's graph as a tree structure.""" def print_tree(node, depth=0): @@ -151,6 +182,10 @@ def print_tree(node, depth=0): for root in root_nodes: print_tree(root) + def show_graph(self, file_name: Optional[str] = None) -> None: + """Display the robot's graph as a directed graph.""" + plot_graph(self.graph, file_name=file_name) + async def _download_assets(self) -> None: """Asynchronously download the assets.""" if not self.assets: @@ -184,6 +219,167 @@ def from_urdf(cls, filename: str) -> "Robot": return robot + @classmethod + def from_url( + cls, name: str, url: str, client: Client, max_depth: int = 0, use_user_defined_root: bool = False + ) -> "Robot": + """Create a robot model from an Onshape CAD assembly.""" + + document = Document.from_url(url=url) + client.set_base_url(document.base_url) + + assembly = client.get_assembly( + did=document.did, + wtype=document.wtype, + wid=document.wid, + eid=document.eid, + log_response=False, + with_meta_data=True, + ) + + instances, occurrences, id_to_name_map = get_instances(assembly=assembly, max_depth=max_depth) + subassemblies, rigid_subassemblies = get_subassemblies(assembly=assembly, client=client, instances=instances) + + parts = get_parts( + assembly=assembly, rigid_subassemblies=rigid_subassemblies, client=client, instances=instances + ) + mates, relations = get_mates_and_relations( + assembly=assembly, + subassemblies=subassemblies, + rigid_subassemblies=rigid_subassemblies, + id_to_name_map=id_to_name_map, + parts=parts, + ) + + graph, root_node = create_graph( + occurrences=occurrences, + instances=instances, + parts=parts, + mates=mates, + use_user_defined_root=use_user_defined_root, + ) + + robot = get_robot( + assembly=assembly, + graph=graph, + root_node=root_node, + parts=parts, + mates=mates, + relations=relations, + client=client, + robot_name=name, + ) + + robot.parts = parts + robot.mates = mates + robot.relations = relations + + robot.subassemblies = subassemblies + robot.rigid_subassemblies = rigid_subassemblies + + robot.assembly = assembly + + return robot + + +def get_robot( + assembly: Assembly, + graph: nx.DiGraph, + root_node: str, + parts: dict[str, Part], + mates: dict[str, MateFeatureData], + relations: dict[str, MateRelationFeatureData], + client: Client, + robot_name: str, +) -> Robot: + """ + Generate a Robot instance from an Onshape assembly. + + Args: + assembly: The Onshape assembly object. + graph: The graph representation of the assembly. + root_node: The root node of the graph. + parts: The dictionary of parts in the assembly. + mates: The dictionary of mates in the assembly. + relations: The dictionary of mate relations in the assembly. + client: The Onshape client object. + robot_name: Name of the robot. + + Returns: + Robot: The generated Robot instance. + """ + robot = Robot(name=robot_name) + + assets_map = {} + stl_to_link_tf_map = {} + topological_mates, topological_relations = get_topological_mates(graph, mates, relations) + + LOGGER.info(f"Processing root node: {root_node}") + + root_link, stl_to_root_tf, root_asset = get_robot_link( + name=root_node, part=parts[root_node], wid=assembly.document.wid, client=client, mate=None + ) + robot.add_link(root_link) + assets_map[root_node] = root_asset + stl_to_link_tf_map[root_node] = stl_to_root_tf + + LOGGER.info(f"Processing {len(graph.edges)} edges in the graph.") + + for parent, child in graph.edges: + mate_key = f"{parent}{MATE_JOINER}{child}" + LOGGER.info(f"Processing edge: {parent} -> {child}") + parent_tf = stl_to_link_tf_map[parent] + + if parent not in parts or child not in parts: + LOGGER.warning(f"Part {parent} or {child} not found in parts dictionary. Skipping.") + continue + + joint_mimic = None + relation = topological_relations.get(topological_mates[mate_key].id) + if relation: + multiplier = ( + relation.relationLength + if relation.relationType == RelationType.RACK_AND_PINION + else relation.relationRatio + ) + joint_mimic = JointMimic( + joint=get_joint_name(relation.mates[RELATION_PARENT].featureId, mates), + multiplier=multiplier, + offset=0.0, + ) + + joint_list, link_list = get_robot_joint( + parent, + child, + topological_mates[mate_key], + parent_tf, + joint_mimic, + is_rigid_assembly=parts[parent].isRigidAssembly, + ) + + link, stl_to_link_tf, asset = get_robot_link( + child, parts[child], assembly.document.wid, client, topological_mates[mate_key] + ) + stl_to_link_tf_map[child] = stl_to_link_tf + assets_map[child] = asset + + if child not in robot.graph: + robot.add_link(link) + else: + LOGGER.warning(f"Link {child} already exists in the robot graph. Skipping.") + + for link in link_list: + if link.name not in robot.graph: + robot.add_link(link) + else: + LOGGER.warning(f"Link {link.name} already exists in the robot graph. Skipping.") + + for joint in joint_list: + robot.add_joint(joint) + + robot.assets = assets_map + return robot + if __name__ == "__main__": LOGGER.set_file_name("robot.log") @@ -203,6 +399,6 @@ def from_urdf(cls, filename: str) -> "Robot": # robot.save() robot = Robot.from_urdf("E:/onshape-api/playground/20240920_umv_mini/20240920_umv_mini/20240920_umv_mini.urdf") - robot.show() - # plot_graph(robot.graph) + # robot.show() + plot_graph(robot.graph) # robot.save(file_path="E:/onshape-api/playground/test.urdf", download_assets=False) diff --git a/onshape_api/urdf.py b/onshape_api/urdf.py index 3c2b0cd..268c643 100644 --- a/onshape_api/urdf.py +++ b/onshape_api/urdf.py @@ -13,12 +13,10 @@ from onshape_api.connect import Asset, Client from onshape_api.log import LOGGER from onshape_api.models.assembly import ( - Assembly, MateFeatureData, MateRelationFeatureData, MateType, Part, - RelationType, ) from onshape_api.models.document import WorkspaceType from onshape_api.models.geometry import MeshGeometry @@ -43,8 +41,7 @@ Origin, VisualLink, ) -from onshape_api.parse import CHILD, MATE_JOINER, PARENT, RELATION_PARENT -from onshape_api.robot import Robot +from onshape_api.parse import CHILD, MATE_JOINER, PARENT from onshape_api.utilities.helpers import get_sanitized_name SCRIPT_DIR = os.path.dirname(__file__) @@ -378,102 +375,3 @@ def get_topological_mates( topological_mates[key] = mates[key] return topological_mates, topological_relations - - -def get_robot( - assembly: Assembly, - graph: DiGraph, - root_node: str, - parts: dict[str, Part], - mates: dict[str, MateFeatureData], - relations: dict[str, MateRelationFeatureData], - client: Client, - robot_name: str, -) -> Robot: - """ - Generate a Robot instance from an Onshape assembly. - - Args: - assembly: The Onshape assembly object. - graph: The graph representation of the assembly. - root_node: The root node of the graph. - parts: The dictionary of parts in the assembly. - mates: The dictionary of mates in the assembly. - relations: The dictionary of mate relations in the assembly. - client: The Onshape client object. - robot_name: Name of the robot. - - Returns: - Robot: The generated Robot instance. - """ - robot = Robot(name=robot_name) - - assets_map = {} - stl_to_link_tf_map = {} - topological_mates, topological_relations = get_topological_mates(graph, mates, relations) - - LOGGER.info(f"Processing root node: {root_node}") - - root_link, stl_to_root_tf, root_asset = get_robot_link( - name=root_node, part=parts[root_node], wid=assembly.document.wid, client=client, mate=None - ) - robot.add_link(root_link) - assets_map[root_node] = root_asset - stl_to_link_tf_map[root_node] = stl_to_root_tf - - LOGGER.info(f"Processing {len(graph.edges)} edges in the graph.") - - for parent, child in graph.edges: - mate_key = f"{parent}{MATE_JOINER}{child}" - LOGGER.info(f"Processing edge: {parent} -> {child}") - parent_tf = stl_to_link_tf_map[parent] - - if parent not in parts or child not in parts: - LOGGER.warning(f"Part {parent} or {child} not found in parts dictionary. Skipping.") - continue - - joint_mimic = None - relation = topological_relations.get(topological_mates[mate_key].id) - if relation: - multiplier = ( - relation.relationLength - if relation.relationType == RelationType.RACK_AND_PINION - else relation.relationRatio - ) - joint_mimic = JointMimic( - joint=get_joint_name(relation.mates[RELATION_PARENT].featureId, mates), - multiplier=multiplier, - offset=0.0, - ) - - joint_list, link_list = get_robot_joint( - parent, - child, - topological_mates[mate_key], - parent_tf, - joint_mimic, - is_rigid_assembly=parts[parent].isRigidAssembly, - ) - - link, stl_to_link_tf, asset = get_robot_link( - child, parts[child], assembly.document.wid, client, topological_mates[mate_key] - ) - stl_to_link_tf_map[child] = stl_to_link_tf - assets_map[child] = asset - - if child not in robot.graph: - robot.add_link(link) - else: - LOGGER.warning(f"Link {child} already exists in the robot graph. Skipping.") - - for link in link_list: - if link.name not in robot.graph: - robot.add_link(link) - else: - LOGGER.warning(f"Link {link.name} already exists in the robot graph. Skipping.") - - for joint in joint_list: - robot.add_joint(joint) - - robot.assets = assets_map - return robot diff --git a/onshape_api/utilities/helpers.py b/onshape_api/utilities/helpers.py index e981108..0e69565 100644 --- a/onshape_api/utilities/helpers.py +++ b/onshape_api/utilities/helpers.py @@ -19,13 +19,25 @@ import matplotlib.animation as animation import matplotlib.pyplot as plt +import numpy as np from PIL import Image from pydantic import BaseModel from onshape_api.log import LOGGER -def save_model_as_json(model: BaseModel, file_path: str) -> None: +class CustomJSONEncoder(json.JSONEncoder): + def default(self, obj): + if isinstance(obj, np.ndarray): + return obj.tolist() # Convert numpy array to list + if isinstance(obj, np.matrix): + return obj.tolist() # Convert numpy matrix to list + if isinstance(obj, set): + return list(obj) # Convert set to list + return super().default(obj) + + +def save_model_as_json(model: BaseModel, file_path: str, indent: int = 4) -> None: """ Save a Pydantic model as a JSON file @@ -45,7 +57,7 @@ def save_model_as_json(model: BaseModel, file_path: str) -> None: """ with open(file_path, "w") as file: - json.dump(model.model_dump(), file, indent=4) + json.dump(model.model_dump(), file, indent=indent, cls=CustomJSONEncoder) def xml_escape(unescaped: str) -> str: