Files
cdsl-cad/onshape_to_cdsl/download_steps.py
T
2026-09-02 13:51:35 +08:00

274 lines
11 KiB
Python

#!/usr/bin/env python3
"""Download AP242 STEP files for public Onshape Part Studios only."""
from __future__ import annotations
import argparse
import base64
import getpass
import hashlib
import json
import os
import re
import ssl
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any
API_VERSION = "v17"
DEFAULT_OUTPUT = Path("onshape_to_cdsl/input/raw")
DEFAULT_SAMPLES = (
(
"00000001",
"https://cad.onshape.com/documents/1ffb81a71e5b402e966b9341/"
"w/6e295017d1b34be684565c40/e/8df255ee6705423d8e85234e",
),
(
"00000002",
"https://cad.onshape.com/documents/1ffb81a71e5b402e966b9341/"
"w/6e295017d1b34be684565c40/e/bb398e4615fe4025b34ea8f0",
),
)
ONSHAPE_URL_RE = re.compile(
r"^https://(?P<stack>[^/]+)/documents/(?P<did>[^/]+)/"
r"(?P<wvm>w|v|m)/(?P<wvmid>[^/]+)/e/(?P<eid>[^/?#]+)"
)
JSON_ACCEPT = "application/json;charset=UTF-8; qs=0.09"
STEP_ACCEPT = "application/step, application/octet-stream"
USER_AGENT = "cdsl-cad-onshape-step-downloader/1.0"
@dataclass(frozen=True)
class PartStudioRef:
sample_id: str
source_url: str
stack: str
did: str
wvm: str
wvmid: str
eid: str
@dataclass
class StepResult:
status: str
path: str | None = None
bytes: int | None = None
sha256: str | None = None
translation_id: str | None = None
http_status: int | None = None
message: str | None = None
def parse_part_studio_url(sample_id: str, url: str) -> PartStudioRef:
match = ONSHAPE_URL_RE.match(url.strip())
if match is None:
raise ValueError(f"not an Onshape Part Studio URL: {url}")
return PartStudioRef(sample_id=sample_id, source_url=url.strip(), **match.groupdict())
def read_url_file(path: Path) -> list[tuple[str, str]]:
"""Read ABC objects YAML lines or plain '<id> <url>' records."""
records: list[tuple[str, str]] = []
url_pattern = re.compile(r"https://cad\.onshape\.com/documents/[^'\"\s]+")
id_pattern = re.compile(r"^\s*['\"]?([A-Za-z0-9_-]+)['\"]?\s*[:\s]")
for line_number, line in enumerate(path.read_text(encoding="utf-8").splitlines(), 1):
url_match = url_pattern.search(line)
if url_match is None:
continue
id_match = id_pattern.search(line)
if id_match is None:
raise ValueError(f"{path}:{line_number}: URL has no leading sample ID")
records.append((id_match.group(1), url_match.group(0)))
if not records:
raise ValueError(f"no Onshape Part Studio URLs found in {path}")
return records
def credentials() -> tuple[str, str]:
access_key = os.environ.get("ONSHAPE_ACCESS_KEY") or getpass.getpass("Onshape access key: ")
secret_key = os.environ.get("ONSHAPE_SECRET_KEY") or getpass.getpass("Onshape secret key: ")
if not access_key or not secret_key:
raise ValueError("both Onshape access and secret keys are required")
return access_key, secret_key
def ca_bundle() -> str | None:
try:
import certifi
except ImportError:
return None
return certifi.where()
class OnshapeClient:
def __init__(self, ref: PartStudioRef, authorization: str, timeout: float) -> None:
self.ref = ref
self.authorization = authorization
self.timeout = timeout
self.context = ssl.create_default_context(cafile=ca_bundle())
def url(self, path: str) -> str:
return f"https://{self.ref.stack}/api/{API_VERSION}{path}"
def request(
self,
method: str,
path: str,
*,
body: dict[str, Any] | None = None,
accept: str = JSON_ACCEPT,
) -> bytes:
payload = None if body is None else json.dumps(body).encode("utf-8")
request = urllib.request.Request(
self.url(path),
data=payload,
method=method,
headers={
"Accept": accept,
"Authorization": self.authorization,
"Content-Type": JSON_ACCEPT,
"User-Agent": USER_AGENT,
},
)
with urllib.request.urlopen(request, timeout=self.timeout, context=self.context) as response:
return response.read()
def get_json(self, path: str) -> Any:
return json.loads(self.request("GET", path).decode("utf-8"))
def sha256(payload: bytes) -> str:
return hashlib.sha256(payload).hexdigest()
def export_step(
client: OnshapeClient,
root: Path,
*,
poll_seconds: float,
poll_limit: int,
) -> StepResult:
ref = client.ref
export_path = f"/partstudios/d/{ref.did}/{ref.wvm}/{ref.wvmid}/e/{ref.eid}/export/step"
body = {
"destinationName": f"{ref.sample_id}.step",
"grouping": True,
"notifyUser": False,
"storeInDocument": False,
"triggerAutoDownload": False,
"stepUnit": "METER",
"stepVersionString": "AP242",
}
try:
translation = json.loads(client.request("POST", export_path, body=body).decode("utf-8"))
translation_id = translation.get("id")
if not isinstance(translation_id, str) or not translation_id:
return StepResult("error", message="STEP export response has no translation id")
translation_path = f"/translations/{urllib.parse.quote(translation_id, safe='')}"
for _ in range(poll_limit):
state = client.get_json(translation_path)
request_state = state.get("requestState") if isinstance(state, dict) else None
if request_state == "DONE":
external_ids = state.get("resultExternalDataIds", [])
if not isinstance(external_ids, list) or not external_ids or not isinstance(external_ids[0], str):
return StepResult("error", translation_id=translation_id, message="completed STEP export has no external data id")
payload = client.request(
"GET",
f"/documents/d/{ref.did}/externaldata/{urllib.parse.quote(external_ids[0], safe='')}",
accept=STEP_ACCEPT,
)
if not payload.startswith(b"ISO-10303-21"):
return StepResult("invalid_step", translation_id=translation_id, message="download is empty or lacks the ISO-10303-21 STEP header")
destination = root / "model.step"
destination.write_bytes(payload)
return StepResult(
"downloaded",
path=destination.name,
bytes=len(payload),
sha256=sha256(payload),
translation_id=translation_id,
)
if request_state == "FAILED":
return StepResult("translation_failed", translation_id=translation_id, message=str(state.get("failureReason") or "translation failed"))
time.sleep(poll_seconds)
return StepResult("translation_timeout", translation_id=translation_id, message=f"STEP translation did not finish after {poll_limit} polls")
except urllib.error.HTTPError as exc:
return StepResult("http_error", http_status=exc.code, message=exc.read().decode("utf-8", "replace")[:1000])
except (urllib.error.URLError, TimeoutError, OSError, ValueError, json.JSONDecodeError) as exc:
return StepResult("error", message=str(exc))
def download_one(ref: PartStudioRef, authorization: str, output: Path, *, timeout: float, poll_seconds: float, poll_limit: int) -> dict[str, Any]:
root = output / ref.sample_id
root.mkdir(parents=True, exist_ok=True)
result = export_step(OnshapeClient(ref, authorization, timeout), root, poll_seconds=poll_seconds, poll_limit=poll_limit)
manifest = {
"schema": "onshape.step_sample.v1",
"source": asdict(ref),
"artifact": {"name": "step", **asdict(result)},
}
(root / "manifest.json").write_text(json.dumps(manifest, ensure_ascii=True, indent=2) + "\n", encoding="utf-8")
return {"sample_id": ref.sample_id, **asdict(result)}
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--url", action="append", default=[], metavar="ID=URL", help="Onshape Part Studio URL; repeatable")
parser.add_argument("--url-file", type=Path, help="ABC objects YAML or text file containing ID and URL records")
parser.add_argument(
"--count",
type=int,
help="maximum samples to download (defaults to the two built-in samples; custom URLs are appended)",
)
parser.add_argument("--out", type=Path, default=DEFAULT_OUTPUT, help="output directory")
parser.add_argument("--timeout", type=float, default=60.0, help="per-request timeout in seconds")
parser.add_argument("--poll-seconds", type=float, default=3.0, help="STEP translation poll interval")
parser.add_argument("--poll-limit", type=int, default=40, help="maximum STEP translation polls")
return parser
def main(argv: list[str] | None = None) -> int:
args = build_parser().parse_args(argv)
if (args.count is not None and args.count < 1) or args.poll_seconds <= 0 or args.poll_limit < 1:
raise ValueError("count and poll limit must be positive; poll interval must be greater than zero")
records: list[tuple[str, str]] = list(DEFAULT_SAMPLES)
for value in args.url:
if "=" not in value:
raise ValueError("--url must use ID=URL syntax")
records.append(tuple(value.split("=", 1)))
if args.url_file:
records.extend(read_url_file(args.url_file))
limit = args.count if args.count is not None else len(records)
if limit < 1:
raise ValueError("count must be positive")
refs = [parse_part_studio_url(sample_id, url) for sample_id, url in records[:limit]]
access_key, secret_key = credentials()
authorization = "Basic " + base64.b64encode(f"{access_key}:{secret_key}".encode("utf-8")).decode("ascii")
results = [download_one(ref, authorization, args.out, timeout=args.timeout, poll_seconds=args.poll_seconds, poll_limit=args.poll_limit) for ref in refs]
summary = {
"schema": "onshape.step_download_summary.v1",
"requested_count": len(refs),
"downloaded_count": sum(item["status"] == "downloaded" for item in results),
"failed_count": sum(item["status"] != "downloaded" for item in results),
"results": results,
}
args.out.mkdir(parents=True, exist_ok=True)
(args.out / "manifest.json").write_text(json.dumps(summary, ensure_ascii=True, indent=2) + "\n", encoding="utf-8")
print(json.dumps(summary, ensure_ascii=True, indent=2))
return 0 if summary["failed_count"] == 0 else 1
if __name__ == "__main__":
try:
raise SystemExit(main())
except (OSError, ValueError) as exc:
print(f"error: {exc}", file=sys.stderr)
raise SystemExit(2)