Added safety rails for parsing functions.

onshape_api/connect.py

@@ -112,9 +112,12 @@ def get_document_metadata(self, did):
         Returns:
             - requests.Response: Onshape response data
         """
-        _request_json = self.request(HTTP.GET, "/api/documents/" + did).json()
+        res = self.request(HTTP.GET, "/api/documents/" + did)
 
-        return DocumentMetaData.model_validate(_request_json)
+        if res.status_code == 404:
+            return None
+
+        return DocumentMetaData.model_validate(res.json())
 
     def get_elements(self, did, wtype, wid):
         """
@@ -236,7 +239,7 @@ def create_assembly(self, did, wid, name="My Assembly"):
 
         return self.request(HTTP.POST, "/api/assemblies/d/" + did + "/w/" + wid, body=payload)
 
-    def get_assembly(self, did, wtype, wid, eid, configuration="default"):
+    def get_assembly(self, did, wtype, wid, eid, configuration="default", log_response=True):
         _request_path = "/api/assemblies/d/" + did + "/" + wtype + "/" + wid + "/e/" + eid
         _assembly_json = self.request(
             HTTP.GET,
@@ -247,6 +250,7 @@ def get_assembly(self, did, wtype, wid, eid, configuration="default"):
                 "includeNonSolids": "false",
                 "configuration": configuration,
             },
+            log_response=log_response,
         ).json()
 
         _assembly = Assembly.model_validate(_assembly_json)
@@ -379,10 +383,13 @@ def _handle_redirect(self, res, method, headers, log_response=True):
         )
 
     def _log_response(self, res):
-        if not 200 <= res.status_code <= 206:
-            LOGGER.debug(f"Request failed, details: {res.text}")
-        else:
-            LOGGER.debug(f"Request succeeded, details: {res.text}")
+        try:
+            if not 200 <= res.status_code <= 206:
+                LOGGER.debug(f"Request failed, details: {res.text}")
+            else:
+                LOGGER.debug(f"Request succeeded, details: {res.text}")
+        except UnicodeEncodeError as e:
+            LOGGER.error(f"UnicodeEncodeError: {e}")
 
     def _make_auth(self, method, date, nonce, path, query=None, ctype="application/json"):
         """

onshape_api/data/preprocess.py

@@ -1,11 +1,15 @@
 import json
 import os
 import re
+from functools import partial
 
 import pandas as pd
 
 from onshape_api.connect import Client
 from onshape_api.models import Assembly
+from onshape_api.models.document import generate_url
+from onshape_api.models.element import ELEMENT_TYPE, Element
+from onshape_api.utilities import LOGGER
 
 AUTOMATE_ASSEMBLYID_PATTERN = r"(?P<documentId>\w{24})_(?P<documentMicroversion>\w{24})_(?P<elementId>\w{24})"
 
@@ -18,53 +22,107 @@ def extract_ids(assembly_id):
         return {"documentId": None, "documentMicroversion": None, "elementId": None}
 
 
-def get_assembly_df(automate_assembly_df):
-    assembly_df = automate_assembly_df["assemblyId"].apply(extract_ids).apply(pd.Series)
+def raise_document_not_exist_error(documentId):
+    raise ValueError(f"Document does not exist: {documentId}")
+
+
+def get_assembly_data(assembly_id: str, client: Client):
+    try:
+        ids = extract_ids(assembly_id)
+        document = client.get_document_metadata(ids["documentId"])
+
+        if document is None:
+            raise_document_not_exist_error(ids["documentId"])
+
+        elements: list[Element] = client.get_elements(
+            did=document.id, wtype=document.defaultWorkspace.type.shorthand, wid=document.defaultWorkspace.id
+        )
+        assembly_ids = [element.id for element in elements.values() if element.elementType == ELEMENT_TYPE.ASSEMBLY]
+
+        ids["elementId"] = assembly_ids
+        ids["wtype"] = document.defaultWorkspace.type.shorthand
+        ids["workspaceId"] = document.defaultWorkspace.id
+
+        LOGGER.info(f"Assembly data retrieved for element: {ids["elementId"]}")
+
+    except Exception as e:
+        LOGGER.warning(f"Error getting assembly data for {assembly_id}")
+        LOGGER.warning(e)
+        ids = {"documentId": None, "documentMicroversion": None, "elementId": None, "wtype": None, "workspaceId": None}
+
+    return ids
+
+
+def get_assembly_df(automate_assembly_df: pd.DataFrame, client: Client) -> pd.DataFrame:
+    """
+    Automate assembly data format:
+        {
+            "assemblyId":"000355ca65fdcbb0e3d825e6_811a5312224ae67ce5b1e180_4bd8ec79e9921e03b989f893_default",
+            "n_subassemblies":1,
+            "n_parts":11,
+            "n_parasolid":11,
+            "n_parasolid_errors":0,
+            "n_step":11,
+            "n_occurrences":11,
+            "n_mates":10,
+            "n_ps_mates":10,
+            "n_step_mates":10,
+            "n_groups":0,
+            "n_relations":0,
+            "is_subassembly":false
+        }
+    """
+    _get_assembly_data = partial(get_assembly_data, client=client)
+    assembly_df = automate_assembly_df["assemblyId"].apply(_get_assembly_data).apply(pd.Series)
     return assembly_df
 
 
 def save_all_jsons(client: Client):
     if not os.path.exists("assemblies.parquet"):
         automate_assembly_df = pd.read_parquet("automate_assemblies.parquet", engine="pyarrow")
-        assembly_df = get_assembly_df(automate_assembly_df)
+        assembly_df = get_assembly_df(automate_assembly_df, client=client)
         assembly_df.to_parquet("assemblies.parquet", engine="pyarrow")
     else:
         assembly_df = pd.read_parquet("assemblies.parquet", engine="pyarrow")
 
-    print(assembly_df.head())
-
-    document = client.get_document_metadata(assembly_df.iloc[0]["documentId"])
-    assembly, assembly_json = client.get_assembly(
-        assembly_df.iloc[0]["documentId"], "w", document.defaultWorkspace.id, assembly_df.iloc[0]["elementId"]
-    )
-
     json_dir = "json"
     os.makedirs(json_dir, exist_ok=True)
 
-    for index, row in assembly_df.iterrows():
+    for _, row in assembly_df.iterrows():
         try:
-            document = client.get_document_metadata(row["documentId"])
-            assembly, assembly_json = client.get_assembly(
-                row["documentId"], "w", document.defaultWorkspace.id, row["elementId"]
-            )
+            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,
+                )
+
+                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']}.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}")
 
-            print(f"Assembly JSON saved to {json_file_path}")
         except Exception as e:
-            print(f"An error occurred for row {index}: {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)
+            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
 
-if __name__ == "__main__":
-    client = Client()
-    # save_all_jsons(client)
 
-    json_file_path = "mate_connectors.json"
-    # json_file_path = "mate_relations.json"
+def validate_assembly_json(json_file_path: str):
     with open(json_file_path) as json_file:
         assembly_json = json.load(json_file)
 
-    assembly = Assembly.model_validate(assembly_json)
-    print(assembly)
+    return Assembly.model_validate(assembly_json)
+
+
+if __name__ == "__main__":
+    client = Client()
+    save_all_jsons(client)

onshape_api/log.py

@@ -116,6 +116,7 @@ def _setup_file_handler(self) -> None:
             mode="w",
             maxBytes=self._file_max_bytes,
             backupCount=self._file_backup_count,
+            encoding="utf-8",
         )
         self._file_handler.setLevel(level=self._file_level.value)
         self._file_handler.setFormatter(fmt=self._std_formatter)

onshape_api/models/document.py

@@ -13,10 +13,37 @@ class WORKSPACE_TYPE(str, Enum):
     M = "m"
 
 
+class META_WORKSPACE_TYPE(str, Enum):
+    WORKSPACE = "workspace"
+    VERSION = "version"
+    MICROVERSION = "microversion"
+
+    @property
+    def shorthand(self) -> str:
+        return self.value[0]
+
+
 DOCUMENT_PATTERN = r"https://cad.onshape.com/documents/([\w\d]+)/(w|v|m)/([\w\d]+)/e/([\w\d]+)"
 
 
+def generate_url(did: str, wtype: str, wid: str, eid: str) -> str:
+    return f"https://cad.onshape.com/documents/{did}/{wtype}/{wid}/e/{eid}"
+
+
 def parse_url(url: str) -> str:
+    """
+    Parse Onshape URL and return document ID, workspace type, workspace ID, and element ID
+
+    Args:
+        url: Onshape URL
+
+    Returns:
+        did: Document ID
+        wtype: Workspace type
+        wid: Workspace ID
+        eid: Element ID
+
+    """
     pattern = re.match(
         DOCUMENT_PATTERN,
         url,
@@ -34,6 +61,18 @@ def parse_url(url: str) -> str:
 
 
 class Document(BaseModel):
+    """
+    Data model for Onshape document
+
+    Attributes:
+        url: Onshape URL
+        did: Document ID
+        wtype: Workspace type
+        wid: Workspace ID
+        eid: Element ID
+
+    """
+
     url: Union[str, None] = None
     did: str
     wtype: str
@@ -43,7 +82,7 @@ class Document(BaseModel):
     def __init__(self, **data):
         super().__init__(**data)
         if self.url is None:
-            self.url = self._generate_url()
+            self.url = generate_url(self.did, self.wtype, self.wid, self.eid)
 
     @field_validator("did", "wid", "eid")
     def check_ids(cls, value: str) -> str:
@@ -70,17 +109,33 @@ def from_url(cls, url: str) -> "Document":
         did, wtype, wid, eid = parse_url(url)
         return cls(url=url, did=did, wtype=wtype, wid=wid, eid=eid)
 
-    def _generate_url(self) -> str:
-        return f"https://cad.onshape.com/documents/{self.did}/{self.wtype}/{self.wid}/e/{self.eid}"
-
 
 class DefaultWorkspace(BaseModel):
+    """
+    Data model for Onshape default workspace
+
+    Attributes:
+        id: Workspace ID
+
+    """
+
     id: str
+    type: META_WORKSPACE_TYPE
 
 
 class DocumentMetaData(BaseModel):
+    """
+    Data model for Onshape document metadata
+
+    Attributes:
+        defaultWorkspace: DefaultWorkspace
+        name: Document name
+
+    """
+
     defaultWorkspace: DefaultWorkspace
     name: str
+    id: str
 
 
 if __name__ == "__main__":

onshape_api/models/element.py

@@ -34,6 +34,9 @@ class ELEMENT_TYPE(str, Enum):
     ASSEMBLY = "ASSEMBLY"
     VARIABLESTUDIO = "VARIABLESTUDIO"
     DRAWING = "DRAWING"
+    BILLOFMATERIALS = "BILLOFMATERIALS"
+    APPLICATION = "APPLICATION"
+    BLOB = "BLOB"
 
 
 class Element(BaseModel):