Fix static type checker errors

Fixes #7

Fix static type checker errors and add type annotations across multiple files.

* **onshape_api/connect.py**
  - Add type annotations to all functions and methods.
  - Fix type errors in `Client` class methods.
  - Ensure all imports are correctly typed.

* **onshape_api/data/preprocess.py**
  - Add type annotations to all functions.
  - Fix type errors in `extract_ids` and `get_assembly_df` functions.

* **onshape_api/graph.py**
  - Add type annotations to all functions.
  - Fix type inconsistencies in `create_graph` and `get_robot_link` functions.

* **onshape_api/log.py**
  - Add type annotations to all methods.
  - Fix type inconsistencies in `Logger` class methods.

* **onshape_api/models/assembly.py**
  - Add type annotations to all classes and methods.
  - Fix type issues in `PartInstance` and `AssemblyInstance` classes.

* **onshape_api/models/document.py**
  - Add type annotations to all classes and methods.
  - Fix type issues in `Document` and `DocumentMetaData` classes.

* **onshape_api/models/element.py**
  - Add type annotations to all classes and methods.
  - Fix type issues in `Element` class.

---

For more details, open the [Copilot Workspace session](https://copilot-workspace.githubnext.com/imsenthur/onshape-api/issues/7?shareId=XXXX-XXXX-XXXX-XXXX).

Author: senthurayyappan

onshape_api/connect.py

@@ -6,7 +6,7 @@
 import secrets
 import string
 from enum import Enum
-from typing import BinaryIO
+from typing import Any, BinaryIO, Optional, Union
 from urllib.parse import parse_qs, urlencode, urlparse
 
 import requests
@@ -32,7 +32,7 @@ class HTTP(str, Enum):
     DELETE = "delete"
 
 
-def load_env_variables(env):
+def load_env_variables(env: str) -> tuple[str, str]:
     """
     Load environment variables from the specified .env file.
 
@@ -61,7 +61,7 @@ def load_env_variables(env):
     return access_key, secret_key
 
 
-def make_nonce():
+def make_nonce() -> str:
     """
     Generate a unique ID for the request, 25 chars in length
 
@@ -85,7 +85,7 @@ class Client:
         - logging (bool, default=True): Turn logging on or off
     """
 
-    def __init__(self, env="./.env", log_file="./onshape.log", log_level=1):
+    def __init__(self, env: str = "./.env", log_file: str = "./onshape.log", log_level: int = 1) -> None:
         """
         Instantiates an instance of the Onshape class. Reads credentials from a .env file.
 
@@ -96,13 +96,13 @@ def __init__(self, env="./.env", log_file="./onshape.log", log_level=1):
             - env (str, default='./.env'): Environment file location
         """
 
-        self._url = BASE_URL
+        self._url: str = BASE_URL
         self._access_key, self._secret_key = load_env_variables(env)
         LOGGER.set_file_name(log_file)
         LOGGER.set_stream_level(LOG_LEVEL[log_level])
         LOGGER.info(f"Onshape API initialized with env file: {env}")
 
-    def get_document(self, did):
+    def get_document(self, did: str) -> DocumentMetaData:
         """
         Get details for a specified document.
 
@@ -116,7 +116,7 @@ def get_document(self, did):
 
         return DocumentMetaData.model_validate(_request_json)
 
-    def get_elements(self, did, wtype, wid):
+    def get_elements(self, did: str, wtype: str, wid: str) -> dict[str, Element]:
         """
         Get list of elements in a document.
 
@@ -138,7 +138,7 @@ def get_elements(self, did, wtype, wid):
 
         return {element["name"]: Element.model_validate(element) for element in _elements_json}
 
-    def get_features_from_partstudio(self, did, wid, eid):
+    def get_features_from_partstudio(self, did: str, wid: str, eid: str) -> requests.Response:
         """
         Gets the feature list for specified document / workspace / part studio.
 
@@ -156,7 +156,7 @@ def get_features_from_partstudio(self, did, wid, eid):
             "/api/partstudios/d/" + did + "/w/" + wid + "/e/" + eid + "/features",
         )
 
-    def get_features_from_assembly(self, did, wtype, wid, eid):
+    def get_features_from_assembly(self, did: str, wtype: str, wid: str, eid: str) -> dict[str, Any]:
         """
         Gets the feature list for specified document / workspace / part studio.
 
@@ -173,7 +173,7 @@ def get_features_from_assembly(self, did, wtype, wid, eid):
             "get", "/api/assemblies/d/" + did + "/" + wtype + "/" + wid + "/e/" + eid + "/features"
         ).json()
 
-    def get_variables(self, did, wid, eid):
+    def get_variables(self, did: str, wid: str, eid: str) -> dict[str, Variable]:
         """
         Get list of variables in a variable studio.
 
@@ -194,7 +194,7 @@ def get_variables(self, did, wid, eid):
 
         return {variable["name"]: Variable.model_validate(variable) for variable in _variables_json[0]["variables"]}
 
-    def set_variables(self, did, wid, eid, variables):
+    def set_variables(self, did: str, wid: str, eid: str, variables: dict[str, Variable]) -> requests.Response:
         """
         Set variables in a variable studio.
 
@@ -219,7 +219,7 @@ def set_variables(self, did, wid, eid, variables):
             body=payload,
         )
 
-    def create_assembly(self, did, wid, name="My Assembly"):
+    def create_assembly(self, did: str, wid: str, name: str = "My Assembly") -> requests.Response:
         """
         Creates a new assembly element in the specified document / workspace.
 
@@ -236,7 +236,9 @@ 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: str, wtype: str, wid: str, eid: str, configuration: str = "default"
+    ) -> tuple[Assembly, dict[str, Any]]:
         _request_path = "/api/assemblies/d/" + did + "/" + wtype + "/" + wid + "/e/" + eid
         _assembly_json = self.request(
             HTTP.GET,
@@ -251,10 +253,10 @@ def get_assembly(self, did, wtype, wid, eid, configuration="default"):
 
         return Assembly.model_validate(_assembly_json), _assembly_json
 
-    def get_parts(self, did, wid, eid):
+    def get_parts(self, did: str, wid: str, eid: str) -> None:
         pass
 
-    def download_stl(self, did, wid, eid, partID, buffer: BinaryIO):
+    def download_stl(self, did: str, wid: str, eid: str, partID: str, buffer: BinaryIO) -> None:
         """
         Exports STL export from a part studio and saves it to a file.
 
@@ -288,7 +290,7 @@ def download_stl(self, did, wid, eid, partID, buffer: BinaryIO):
         else:
             LOGGER.info(f"Failed to download STL file: {response.status_code} - {response.text}")
 
-    def get_mass_properties(self, did, wid, eid, partID):
+    def get_mass_properties(self, did: str, wid: str, eid: str, partID: str) -> MassModel:
         """
         Get mass properties for a part in a part studio.
 
@@ -306,7 +308,16 @@ def get_mass_properties(self, did, wid, eid, partID):
 
         return MassModel.model_validate(_resonse_json["bodies"][partID])
 
-    def request(self, method, path, query=None, headers=None, body=None, base_url=None, log_response=True):
+    def request(
+        self,
+        method: Union[HTTP, str],
+        path: str,
+        query: Optional[dict[str, Any]] = None,
+        headers: Optional[dict[str, str]] = None,
+        body: Optional[dict[str, Any]] = None,
+        base_url: Optional[str] = None,
+        log_response: bool = True,
+    ) -> requests.Response:
         """
         Issues a request to Onshape
 
@@ -347,10 +358,12 @@ def request(self, method, path, query=None, headers=None, body=None, base_url=No
 
         return res
 
-    def _build_url(self, base_url, path, query):
+    def _build_url(self, base_url: str, path: str, query: dict[str, Any]) -> str:
         return base_url + path + "?" + urlencode(query)
 
-    def _send_request(self, method, url, headers, body):
+    def _send_request(
+        self, method: Union[HTTP, str], url: str, headers: dict[str, str], body: Optional[dict[str, Any]]
+    ) -> requests.Response:
         return requests.request(
             method,
             url,
@@ -361,7 +374,9 @@ def _send_request(self, method, url, headers, body):
             timeout=10,  # Specify an appropriate timeout value in seconds
         )
 
-    def _handle_redirect(self, res, method, headers, log_response=True):
+    def _handle_redirect(
+        self, res: requests.Response, method: Union[HTTP, str], headers: dict[str, str], log_response: bool = True
+    ) -> requests.Response:
         location = urlparse(res.headers["Location"])
         querystring = parse_qs(location.query)
 
@@ -374,13 +389,21 @@ def _handle_redirect(self, res, method, headers, log_response=True):
             method, location.path, query=new_query, headers=headers, base_url=new_base_url, log_response=log_response
         )
 
-    def _log_response(self, res):
+    def _log_response(self, res: requests.Response) -> None:
         if not 200 <= res.status_code <= 206:
             LOGGER.debug(f"Request failed, details: {res.text}")
         else:
             LOGGER.debug(f"Request succeeded, details: {res.text}")
 
-    def _make_auth(self, method, date, nonce, path, query=None, ctype="application/json"):
+    def _make_auth(
+        self,
+        method: Union[HTTP, str],
+        date: str,
+        nonce: str,
+        path: str,
+        query: Optional[dict[str, Any]] = None,
+        ctype: str = "application/json",
+    ) -> str:
         """
         Create the request signature to authenticate
 
@@ -412,7 +435,13 @@ def _make_auth(self, method, date, nonce, path, query=None, ctype="application/j
 
         return auth
 
-    def _make_headers(self, method, path, query=None, headers=None):
+    def _make_headers(
+        self,
+        method: Union[HTTP, str],
+        path: str,
+        query: Optional[dict[str, Any]] = None,
+        headers: Optional[dict[str, str]] = None,
+    ) -> dict[str, str]:
         """
         Creates a headers object to sign the request
 

onshape_api/data/preprocess.py

@@ -1,27 +1,22 @@
 import json
 import os
 import re
-
 import pandas as pd
-
 import onshape_api as osa
 
 AUTOMATE_ASSEMBLYID_PATTERN = r"(?P<documentId>\w{24})_(?P<documentMicroversion>\w{24})_(?P<elementId>\w{24})"
 
-
-def extract_ids(assembly_id):
+def extract_ids(assembly_id: str) -> dict[str, str | None]:
     match = re.match(AUTOMATE_ASSEMBLYID_PATTERN, assembly_id)
     if match:
         return match.groupdict()
     else:
         return {"documentId": None, "documentMicroversion": None, "elementId": None}
 
-
-def get_assembly_df(automate_assembly_df):
+def get_assembly_df(automate_assembly_df: pd.DataFrame) -> pd.DataFrame:
     assembly_df = automate_assembly_df["assemblyId"].apply(extract_ids).apply(pd.Series)
     return assembly_df
 
-
 if __name__ == "__main__":
     client = osa.Client()
 

onshape_api/graph.py

@@ -1,7 +1,7 @@
 import io
 import os
 import random
-from typing import Optional, Union
+from typing import Optional, Union, Dict, Tuple, List
 
 import matplotlib.pyplot as plt
 import networkx as nx
@@ -39,7 +39,7 @@
 CURRENT_DIR = os.getcwd()
 
 
-def generate_names(max_length: int) -> list[str]:
+def generate_names(max_length: int) -> List[str]:
     words_file_path = os.path.join(SCRIPT_DIR, "words.txt")
 
     with open(words_file_path) as file:
@@ -51,25 +51,25 @@ def generate_names(max_length: int) -> list[str]:
     return random.sample(words, max_length)
 
 
-def show_graph(graph: nx.Graph):
+def show_graph(graph: nx.Graph) -> None:
     nx.draw_circular(graph, with_labels=True)
     plt.show()
 
 
-def convert_to_digraph(graph: nx.Graph) -> nx.DiGraph:
+def convert_to_digraph(graph: nx.Graph) -> Tuple[nx.DiGraph, str]:
     _centrality = nx.closeness_centrality(graph)
     _root_node = max(_centrality, key=_centrality.get)
     _graph = nx.bfs_tree(graph, _root_node)
     return _graph, _root_node
 
 
 def create_graph(
-    occurences: dict[str, Occurrence],
-    instances: dict[str, Instance],
-    parts: dict[str, Part],
-    mates: dict[str, MateFeatureData],
+    occurences: Dict[str, Occurrence],
+    instances: Dict[str, Instance],
+    parts: Dict[str, Part],
+    mates: Dict[str, MateFeatureData],
     directed: bool = True,
-):
+) -> Union[nx.Graph, Tuple[nx.DiGraph, str]]:
     graph = nx.Graph()
 
     for occurence in occurences:
@@ -90,14 +90,14 @@ def create_graph(
             LOGGER.warning(f"Mate {mate} not found")
 
     if directed:
-        graph = convert_to_digraph(graph)
+        graph, root_node = convert_to_digraph(graph)
 
     LOGGER.info(f"Graph created with {len(graph.nodes)} nodes and {len(graph.edges)} edges")
 
     return graph
 
 
-def download_stl_mesh(did, wid, eid, partID, client: Client, transform: np.ndarray, file_name: str) -> str:
+def download_stl_mesh(did: str, wid: str, eid: str, partID: str, client: Client, transform: np.ndarray, file_name: str) -> str:
     try:
         with io.BytesIO() as buffer:
             LOGGER.info(f"Downloading mesh for {file_name}...")
@@ -128,7 +128,7 @@ def get_robot_link(
     workspaceId: str,
     client: Client,
     mate: Optional[Union[MateFeatureData, None]] = None,
-):
+) -> Tuple[Link, np.matrix]:
     LOGGER.info(f"Creating robot link for {name}")
 
     if mate is None:
@@ -190,7 +190,7 @@ def get_robot_joint(
     child: str,
     mate: MateFeatureData,
     stl_to_parent_tf: np.matrix,
-):
+) -> Union[RevoluteJoint, FixedJoint]:
     LOGGER.info(f"Creating robot joint from {parent} to {child}")
 
     parent_to_mate_tf = mate.matedEntities[1].matedCS.part_to_mate_tf
@@ -228,20 +228,20 @@ def get_robot_joint(
 def get_urdf_components(
     graph: Union[nx.Graph, nx.DiGraph],
     workspaceId: str,
-    parts: dict[str, Part],
-    mass_properties: dict[str, MassModel],
-    mates: dict[str, MateFeatureData],
+    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):
         graph, root_node = convert_to_digraph(graph)
 
-    joints = []
-    links = []
+    joints: List[Union[RevoluteJoint, FixedJoint]] = []
+    links: List[Link] = []
 
     _readable_names = generate_names(len(graph.nodes))
     _readable_names_mapping = dict(zip(graph.nodes, _readable_names))
-    _stl_to_link_tf_mapping = {}
+    _stl_to_link_tf_mapping: Dict[str, np.matrix] = {}
 
     LOGGER.info(f"Processing root node: {_readable_names_mapping[root_node]}")