Restructured convenience methods and created from_url and from_urdf c…

…lass methods within Robot class.

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()

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()

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:
     """

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,15 +332,15 @@ 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)
 
     tasks = []
     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,15 +489,15 @@ 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))
                     parts[parent_occurrences[0]].rigidAssemblyToPartTF[parent_occurrences[1]] = parent_parentCS.part_tf
                     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,30 +526,30 @@ 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]]:
     """
     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(
         assembly.rootAssembly.features,
         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,14 +559,14 @@ 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]]:
     """
     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)
     )