Enforcing CRLF and minor ruff fixes.

examples/bike/main.py

@@ -21,10 +21,10 @@
 
 variables["wheelDiameter"].expression = "300 mm"
 variables["wheelThickness"].expression = "71 mm"
-variables["forkAngle"].expression = "30 deg"
+variables["forkAngle"].expression = "20 deg"
 
 client.set_variables(doc.did, doc.wid, elements["variables"].id, variables)
-assembly = client.get_assembly(doc.did, doc.wtype, doc.wid, elements["assembly"].id)
+assembly, _ = client.get_assembly(doc.did, doc.wtype, doc.wid, elements["assembly"].id)
 
 occurences = get_occurences(assembly)
 instances = get_instances(assembly)

onshape_api/graph.py

@@ -7,6 +7,7 @@
 import networkx as nx
 import numpy as np
 import stl
+from networkx import DiGraph, Graph
 
 from onshape_api.connect import Client
 from onshape_api.log import LOGGER
@@ -51,12 +52,16 @@ def generate_names(max_length: int) -> list[str]:
     return random.sample(words, max_length)
 
 
-def show_graph(graph: nx.Graph) -> None:
+def get_random_color() -> tuple[float, float, float, float]:
+    return random.SystemRandom().choice(list(COLORS)).value
+
+
+def show_graph(graph: Graph) -> None:
     nx.draw_circular(graph, with_labels=True)
     plt.show()
 
 
-def convert_to_digraph(graph: nx.Graph) -> tuple[nx.DiGraph, str]:
+def convert_to_digraph(graph: Graph) -> tuple[DiGraph, str]:
     _centrality = nx.closeness_centrality(graph)
     _root_node = max(_centrality, key=_centrality.get)
     _graph = nx.bfs_tree(graph, _root_node)
@@ -69,8 +74,8 @@ def create_graph(
     parts: dict[str, Part],
     mates: dict[str, MateFeatureData],
     directed: bool = True,
-) -> Union[nx.Graph, tuple[nx.DiGraph, str]]:
-    graph = nx.Graph()
+) -> Union[Graph, tuple[DiGraph, str]]:
+    graph: Graph = Graph()
 
     for occurence in occurences:
         if instances[occurence].type == InstanceType.PART:
@@ -161,7 +166,7 @@ def get_robot_link(
         visual=VisualLink(
             origin=_origin,
             geometry=MeshGeometry(_mesh_path),
-            material=Material.from_color(name=f"{name}_material", color=random.SystemRandom().choice(list(COLORS))),
+            material=Material.from_color(name=f"{name}_material", color=get_random_color()),
         ),
         inertial=InertialLink(
             origin=Origin(
@@ -228,14 +233,14 @@ def get_robot_joint(
 
 
 def get_urdf_components(
-    graph: Union[nx.Graph, nx.DiGraph],
+    graph: Union[Graph, DiGraph],
     workspaceId: str,
     parts: dict[str, Part],
     mass_properties: dict[str, MassModel],
     mates: dict[str, MateFeatureData],
     client: Client,
 ) -> tuple[list[Link], list[Union[RevoluteJoint, FixedJoint]]]:
-    if not isinstance(graph, nx.DiGraph):
+    if not isinstance(graph, DiGraph):
         graph, root_node = convert_to_digraph(graph)
 
     joints: list[Union[RevoluteJoint, FixedJoint]] = []
@@ -287,3 +292,7 @@ def get_urdf_components(
         links.append(_link)
 
     return links, joints
+
+
+if __name__ == "__main__":
+    print(get_random_color())

onshape_api/models/link.py

@@ -102,7 +102,7 @@ def to_xml(self, root: ET.Element | None = None) -> ET.Element:
         return material
 
     @classmethod
-    def from_color(cls, name: str, color: COLORS) -> "Material":
+    def from_color(cls, name: str, color: tuple[float, float, float, float]) -> "Material":
         return cls(name, color)
 
 

onshape_api/models/mass.py

@@ -90,13 +90,13 @@ def inertia_wrt(self, reference: np.matrix) -> np.matrix:
 
         return reference @ self.inertia_matrix @ reference.T
 
-    def center_of_mass_wrt(self, reference: np.matrix) -> np.ndarray:
+    def center_of_mass_wrt(self, reference: np.matrix) -> tuple[float, float, float]:
         if reference.shape != (4, 4):
             raise ValueError("Reference frame must be a 4x4 matrix")
 
         com = np.matrix([*list(self.center_of_mass), 1.0])
         com_wrt = (reference * com.T)[:3]
-        return np.array([com_wrt[0, 0], com_wrt[1, 0], com_wrt[2, 0]])
+        return (com_wrt[0, 0], com_wrt[1, 0], com_wrt[2, 0])
 
 
 if __name__ == "__main__":

pyproject.toml

@@ -30,6 +30,7 @@ deptry = "^0.16.2"
 mypy = "^1.5.1"
 pre-commit = "^3.4.0"
 tox = "^4.11.1"
+types-networkx = "^3.3.0.20241020"
 
 [tool.poetry.group.docs.dependencies]
 mkdocs = "^1.4.2"