Caught some exceptions within the get_document_metadata method and ad…

…ded more safety rails to assembly validations.

onshape_api/connect.py

@@ -111,10 +111,28 @@ def get_document_metadata(self, did):
 
         Returns:
             - requests.Response: Onshape response data
+
         """
         res = self.request(HTTP.GET, "/api/documents/" + did)
 
-        if res.status_code == 404:
+        if res.status_code == 404 or res.status_code == 403:
+            """
+            404: Document not found
+                {
+                    "message": "Not found.",
+                    "code": 0,
+                    "status": 404,
+                    "moreInfoUrl": ""
+                }
+
+            403: Resource does not exist
+                {
+                    "message": "Resource does not exist, or you do not have permission to access it.",
+                    "code": 1002,
+                    "status": 403,
+                    "moreInfoUrl": null
+                }
+            """
             return None
 
         return DocumentMetaData.model_validate(res.json())

onshape_api/data/preprocess.py

@@ -43,7 +43,7 @@ def get_assembly_data(assembly_id: str, client: Client):
         ids["wtype"] = document.defaultWorkspace.type.shorthand
         ids["workspaceId"] = document.defaultWorkspace.id
 
-        LOGGER.info(f"Assembly data retrieved for element: {ids["elementId"]}")
+        LOGGER.info(f"Assembly data retrieved for element: {ids['elementId']}")
 
     except Exception as e:
         LOGGER.warning(f"Error getting assembly data for {assembly_id}")

onshape_api/graph.py

@@ -1,13 +1,16 @@
+from typing import Union
+
 import matplotlib.pyplot as plt
 import networkx as nx
 
 from onshape_api.log import LOGGER
 from onshape_api.models.assembly import (
-    Instance,
+    AssemblyInstance,
     InstanceType,
     MateFeatureData,
     Occurrence,
     Part,
+    PartInstance,
 )
 from onshape_api.parse import MATE_JOINER
 
@@ -31,7 +34,7 @@ def convert_to_digraph(graph: nx.Graph) -> nx.DiGraph:
 
 def create_graph(
     occurences: dict[str, Occurrence],
-    instances: dict[str, Instance],
+    instances: dict[str, Union[PartInstance, AssemblyInstance]],
     parts: dict[str, Part],
     mates: dict[str, MateFeatureData],
     directed: bool = True,

onshape_api/models/assembly.py

@@ -1,20 +1,77 @@
+"""
+This module contains the data models for the Assembly API responses from Onshape REST API. The data models are Pydantic
+BaseModel classes that are used to parse the JSON responses from the API into Python objects. The data models are used
+to validate the JSON responses and to provide type hints for the data structures.
+
+Models:
+- Occurrence: Occurence model
+- Part: Part data model
+- PartInstance: Part Instance model
+- AssemblyInstance: Assembly Instance model
+- AssemblyFeature: AssemblyFeature data model
+- Pattern: Pattern data model
+- SubAssembly: SubAssembly data model
+- RootAssembly: RootAssembly data model
+- Assembly: Assembly data model
+
+Supplementary models:
+- IDBase: Base model for Part, SubAssembly, and AssemblyInstance in Assembly context
+- MatedCS: Mated CS model
+- MatedEntity: MatedEntity data model
+- MateRelationMate: MateRelationMate data model
+- MateGroupFeatureOccurrence: MateGroupFeatureOccurrence data model
+- MateGroupFeatureData: MateGroupFeatureData data model
+- MateConnectorFeatureData: MateConnectorFeatureData data model
+- MateRelationFeatureData: MateRelationFeatureData data model
+- MateFeatureData: MateFeatureData data model
+
+Enums:
+- INSTANCE_TYPE: Instance type to distinguish between Part and Assembly
+- MATE_TYPE: Type of mate between two parts or assemblies, e.g. SLIDER, CYLINDRICAL, REVOLUTE, etc.
+- RELATION_TYPE: Type of mate relation between two parts or assemblies, e.g. LINEAR, GEAR, SCREW, etc.
+- ASSEMBLY_FEATURE_TYPE: Type of assembly feature, e.g. mate, mateRelation, mateGroup, mateConnector
+
+"""
+
 from enum import Enum
 from typing import Union
 
 import numpy as np
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, Field, field_validator
 
 from onshape_api.models.document import Document
 from onshape_api.models.mass import MassProperties
 from onshape_api.utilities.helpers import generate_uid
 
 
 class INSTANCE_TYPE(str, Enum):
+    """
+    Enum to distinguish between Part and Assembly.
+
+    Attributes:
+        PART (str): Represents a part instance.
+        ASSEMBLY (str): Represents an assembly instance.
+    """
+
     PART = "Part"
     ASSEMBLY = "Assembly"
 
 
 class MATE_TYPE(str, Enum):
+    """
+    Enum to represent the type of mate between two parts or assemblies.
+
+    Attributes:
+        SLIDER (str): Represents a slider mate.
+        CYLINDRICAL (str): Represents a cylindrical mate.
+        REVOLUTE (str): Represents a revolute mate.
+        PIN_SLOT (str): Represents a pin-slot mate.
+        PLANAR (str): Represents a planar mate.
+        BALL (str): Represents a ball mate.
+        FASTENED (str): Represents a fastened mate.
+        PARALLEL (str): Represents a parallel mate.
+    """
+
     SLIDER = "SLIDER"
     CYLINDRICAL = "CYLINDRICAL"
     REVOLUTE = "REVOLUTE"
@@ -26,13 +83,33 @@ class MATE_TYPE(str, Enum):
 
 
 class RELATION_TYPE(str, Enum):
+    """
+    Enum to represent the type of mate relation between two parts or assemblies.
+
+    Attributes:
+        LINEAR (str): Represents a linear relation.
+        GEAR (str): Represents a gear relation.
+        SCREW (str): Represents a screw relation.
+        RACK_AND_PINION (str): Represents a rack and pinion relation.
+    """
+
     LINEAR = "LINEAR"
     GEAR = "GEAR"
     SCREW = "SCREW"
     RACK_AND_PINION = "RACK_AND_PINION"
 
 
 class ASSEMBLY_FEATURE_TYPE(str, Enum):
+    """
+    Enum to represent the type of assembly feature.
+
+    Attributes:
+        MATE (str): Represents a mate feature.
+        MATERELATION (str): Represents a mate relation feature.
+        MATEGROUP (str): Represents a mate group feature.
+        MATECONNECTOR (str): Represents a mate connector feature.
+    """
+
     MATE = "mate"
     MATERELATION = "mateRelation"
     MATEGROUP = "mateGroup"
@@ -41,27 +118,47 @@ class ASSEMBLY_FEATURE_TYPE(str, Enum):
 
 class Occurrence(BaseModel):
     """
-    Occurence model
+    Occurrence model representing the state of an instance in an assembly.
 
+    Example JSON representation:
     {
-    "fixed" : false,
-    "transform" :
-        [ 0.8660254037844396, 0.0, 0.5000000000000004, 0.09583333333333346,
-        0.0, 1.0, 0.0, -1.53080849893419E-19,
-        -0.5000000000000004, 0.0, 0.8660254037844396, 0.16598820239201767,
-        0.0, 0.0, 0.0, 1.0 ],
-    "hidden" : false,
-    "path" : [ "M0Cyvy+yIq8Rd7En0" ]
+        "fixed": false,
+        "transform": [
+            0.8660254037844396, 0.0, 0.5000000000000004, 0.09583333333333346,
+            0.0, 1.0, 0.0, -1.53080849893419E-19,
+            -0.5000000000000004, 0.0, 0.8660254037844396, 0.16598820239201767,
+            0.0, 0.0, 0.0, 1.0
+        ],
+        "hidden": false,
+        "path": ["M0Cyvy+yIq8Rd7En0"]
     }
+
+    Attributes:
+        fixed (bool): Indicates if the occurrence is fixed in space.
+        transform (list[float]): A 4x4 transformation matrix represented as a list of 16 floats.
+        hidden (bool): Indicates if the occurrence is hidden.
+        path (list[str]): A list of strings representing the path to the instance.
     """
 
-    fixed: bool
-    transform: list[float]
-    hidden: bool
-    path: list[str]
+    fixed: bool = Field(..., description="Indicates if the occurrence is fixed in space.")
+    transform: list[float] = Field(..., description="A 4x4 transformation matrix represented as a list of 16 floats.")
+    hidden: bool = Field(..., description="Indicates if the occurrence is hidden.")
+    path: list[str] = Field(..., description="A list of strings representing the path to the instance.")
 
     @field_validator("transform")
     def check_transform(cls, v: list[float]) -> list[float]:
+        """
+        Validates that the transform list has exactly 16 values.
+
+        Args:
+            v (list[float]): The transform list to validate.
+
+        Returns:
+            list[float]: The validated transform list.
+
+        Raises:
+            ValueError: If the transform list does not contain exactly 16 values.
+        """
         if len(v) != 16:
             raise ValueError("Transform must have 16 values")
 
@@ -70,57 +167,101 @@ def check_transform(cls, v: list[float]) -> list[float]:
 
 class IDBase(BaseModel):
     """
-    Base model for Part in Assembly context
+    Base model for Part, SubAssembly, and AssemblyInstance in Assembly context.
+
+    Example JSON representation:
     {
         "fullConfiguration" : "default",
         "configuration" : "default",
         "documentId" : "a1c1addf75444f54b504f25c",
         "elementId" : "0b0c209535554345432581fe",
         "documentMicroversion" : "12fabf866bef5a9114d8c4d2"
     }
+
+    Attributes:
+        fullConfiguration (str): The full configuration of the entity.
+        configuration (str): The configuration of the entity.
+        documentId (str): The document ID of the entity.
+        elementId (str): The element ID of the entity.
+        documentMicroversion (str): The microversion of the document.
     """
 
-    fullConfiguration: str
-    configuration: str
-    documentId: str
-    elementId: str
-    documentMicroversion: str
+    fullConfiguration: str = Field(..., description="The full configuration of the entity.")
+    configuration: str = Field(..., description="The configuration of the entity.")
+    documentId: str = Field(..., description="The document ID of the entity.")
+    elementId: str = Field(..., description="The element ID of the entity.")
+    documentMicroversion: str = Field(..., description="The microversion of the document.")
 
     @field_validator("documentId", "elementId", "documentMicroversion")
     def check_ids(cls, v: str) -> str:
+        """
+        Validates that the ID fields have exactly 24 characters.
+
+        Args:
+            v (str): The ID field to validate.
+
+        Returns:
+            str: The validated ID field.
+
+        Raises:
+            ValueError: If the ID field does not contain exactly 24 characters.
+        """
         if len(v) != 24:
             raise ValueError("DocumentId must have 24 characters")
 
         return v
 
     @property
     def uid(self) -> str:
+        """
+        Generates a unique identifier for the part.
+
+        Returns:
+            str: The unique identifier generated from documentId, documentMicroversion,
+                elementId, and fullConfiguration.
+        """
         return generate_uid([self.documentId, self.documentMicroversion, self.elementId, self.fullConfiguration])
 
 
 class Part(IDBase):
     """
-    Part data model
+    Part data model representing a part in an assembly.
+
+    Example JSON representation:
     {
-      "isStandardContent" : false,
-      "partId" : "RDBD",
-      "bodyType" : "solid",
-
-      "fullConfiguration" : "default",
-      "configuration" : "default",
-      "documentId" : "a1c1addf75444f54b504f25c",
-      "elementId" : "0b0c209535554345432581fe",
-      "documentMicroversion" : "349f6413cafefe8fb4ab3b07"
+        "isStandardContent": false,
+        "partId": "RDBD",
+        "bodyType": "solid",
+        "fullConfiguration": "default",
+        "configuration": "default",
+        "documentId": "a1c1addf75444f54b504f25c",
+        "elementId": "0b0c209535554345432581fe",
+        "documentMicroversion": "349f6413cafefe8fb4ab3b07"
     }
+
+    Attributes:
+        isStandardContent (bool): Indicates if the part is standard content.
+        partId (str): The unique identifier of the part.
+        bodyType (str): The type of the body (e.g., solid, surface).
+        MassProperty (Union[MassProperties, None]): The mass properties of the part, if available.
     """
 
-    isStandardContent: bool
-    partId: str
-    bodyType: str
-    MassProperty: Union[MassProperties, None] = None
+    isStandardContent: bool = Field(..., description="Indicates if the part is standard content.")
+    partId: str = Field(..., description="The unique identifier of the part.")
+    bodyType: str = Field(..., description="The type of the body (e.g., solid, surface).")
+    MassProperty: Union[MassProperties, None] = Field(
+        None, description="The mass properties of the part, if available."
+    )
 
     @property
     def uid(self) -> str:
+        """
+        Generates a unique identifier for the part.
+
+        Returns:
+            str: The unique identifier generated from documentId, documentMicroversion,
+                elementId, partId, and fullConfiguration.
+        """
         return generate_uid([
             self.documentId,
             self.documentMicroversion,
@@ -204,9 +345,6 @@ def check_type(cls, v: INSTANCE_TYPE) -> INSTANCE_TYPE:
         return v
 
 
-Instance = Union[PartInstance, AssemblyInstance]
-
-
 class MatedCS(BaseModel):
     """
     Mated CS model

onshape_api/models/element.py

@@ -37,6 +37,7 @@ class ELEMENT_TYPE(str, Enum):
     BILLOFMATERIALS = "BILLOFMATERIALS"
     APPLICATION = "APPLICATION"
     BLOB = "BLOB"
+    FEATURESTUDIO = "FEATURESTUDIO"
 
 
 class Element(BaseModel):