refactor init

kol/geometry.py

@@ -38,6 +38,10 @@ def transform(points: np.ndarray) -> np.ndarray:
     mesh.normals = transform(mesh.normals)
 
 
+def b_to_a_tf(a_to_b_tf: np.matrix) -> np.matrix:
+    return np.matrix(np.linalg.inv(a_to_b_tf))
+
+
 @dataclass
 class Dynamics:
     mass: float

kol/onshape/api.py

@@ -2,9 +2,8 @@
 
 import logging
 from pathlib import Path
-from typing import Literal
 
-from kol.onshape.client import OnshapeClient
+from kol.onshape.client import DocumentInfo, OnshapeClient, WorkspaceType
 from kol.onshape.schema.assembly import Assembly, AssemblyMetadata, Part, RootAssembly, SubAssembly
 from kol.onshape.schema.document import Document
 from kol.onshape.schema.elements import Elements, ElementType
@@ -18,17 +17,15 @@ def escape_url(s: str) -> str:
     return s.replace("/", "%2f").replace("+", "%2b")
 
 
-DEFAULT_BASE_URL = "https://cad.onshape.com"
-
-WorkspaceType = Literal["w", "v"]
-
-
 class OnshapeApi:
     def __init__(self, client: OnshapeClient) -> None:
         super().__init__()
 
         self.client = client
 
+    def parse_url(self, document_url: str) -> DocumentInfo:
+        return self.client.parse_url(document_url)
+
     def get_document(self, did: str) -> Document:
         data = self.client.request("get", f"/api/documents/{did}").json()
         return Document.model_validate(data)
@@ -50,15 +47,11 @@ def get_first_assembly_id(self, document_id: str, workspace_id: str, workspace_t
                 return element.id
         raise ValueError("Assembly not found")
 
-    def get_assembly(
-        self,
-        document_id: str,
-        workspace_id: str,
-        element_id: str,
-        workspace_type: WorkspaceType = "w",
-        configuration: str = "default",
-    ) -> Assembly:
-        path = f"/api/assemblies/d/{document_id}/{workspace_type}/{workspace_id}/e/{element_id}"
+    def get_assembly(self, document: DocumentInfo, configuration: str = "default") -> Assembly:
+        path = (
+            f"/api/assemblies/d/{document.document_id}/"
+            f"{document.item_kind}/{document.item_id}/e/{document.element_id}"
+        )
         data = self.client.request(
             "get",
             path,

kol/onshape/client.py

@@ -8,9 +8,11 @@
 import logging
 import os
 import random
+import re
 import string
 import urllib.parse
-from typing import Any, Literal
+from dataclasses import dataclass
+from typing import Any, Literal, cast
 
 import requests
 
@@ -28,6 +30,17 @@ def escape_url(s: str) -> str:
 Method = Literal["get", "post", "put", "delete"]
 
 
+WorkspaceType = Literal["w", "v"]
+
+
+@dataclass
+class DocumentInfo:
+    document_id: str
+    item_kind: WorkspaceType
+    item_id: str
+    element_id: str
+
+
 class OnshapeClient:
     def __init__(
         self,
@@ -59,6 +72,16 @@ def __init__(
         self.secret_key = secret_key.encode("utf-8")
         self.base_url = base_url
 
+    def parse_url(self, document_url: str) -> DocumentInfo:
+        url_match = re.match(rf"{self.base_url}/documents/([\w\d]+)/(w|v)/([\w\d]+)/e/([\w\d]+)", document_url)
+        if url_match is None:
+            raise ValueError(f"Invalid document URL: {document_url}")
+        document_id = url_match.group(1)
+        item_kind = cast(WorkspaceType, url_match.group(2))
+        item_id = url_match.group(3)
+        element_id = url_match.group(4)
+        return DocumentInfo(document_id, item_kind, item_id, element_id)
+
     def _make_nonce(self) -> str:
         """Generate a unique ID for the request, 25 chars in length.