From 77550cfd67a979f750d568800da437c2337d0c97 Mon Sep 17 00:00:00 2001 From: imsenthur Date: Sat, 16 Nov 2024 15:21:30 -0500 Subject: [PATCH] Added support for assembly names. --- benchmark/assembly.py | 74 ++++++++++++++++----------------- onshape_api/connect.py | 75 +++++++++++++++++++++++++++++++++- onshape_api/models/assembly.py | 2 + onshape_api/parse.py | 6 ++- 4 files changed, 117 insertions(+), 40 deletions(-) diff --git a/benchmark/assembly.py b/benchmark/assembly.py index e5d3cd4..c498f4c 100644 --- a/benchmark/assembly.py +++ b/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") diff --git a/onshape_api/connect.py b/onshape_api/connect.py index 31b39f3..2d6f59a 100644 --- a/onshape_api/connect.py +++ b/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: diff --git a/onshape_api/models/assembly.py b/onshape_api/models/assembly.py index 539dc7d..959399b 100644 --- a/onshape_api/models/assembly.py +++ b/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.") diff --git a/onshape_api/parse.py b/onshape_api/parse.py index faa34e3..442678c 100644 --- a/onshape_api/parse.py +++ b/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(...)