kscale 0.0.2__py3-none-any.whl → 0.0.5__py3-none-any.whl

kscale/__init__.py

@@ -1 +1,5 @@
-__version__ = "0.0.2"
+"""Defines the common interface for the K-Scale Python API."""
+
+__version__ = "0.0.5"
+
+from kscale.api import KScale

kscale/api.py

@@ -0,0 +1,14 @@
+"""Defines common functionality for the K-Scale API."""
+
+from kscale.store.api import StoreAPI
+from kscale.utils.api_base import APIBase
+
+
+class KScale(
+    StoreAPI,
+    APIBase,
+):
+    """Defines a common interface for the K-Scale API."""
+
+    def __init__(self, api_key: str | None = None) -> None:
+        self.api_key = api_key

kscale/conf.py

@@ -17,7 +17,7 @@ def get_path() -> Path:
 
 @dataclass
 class StoreSettings:
-    api_key: str = field(default=II("oc.env:KSCALE_API_KEY,"))
+    api_key: str | None = field(default=None)
     cache_dir: str = field(default=II("oc.env:KSCALE_CACHE_DIR,'~/.kscale/cache/'"))
 
 

kscale/store/__init__.py

@@ -1 +0,0 @@
-__version__ = "0.0.1"

kscale/store/api.py

@@ -0,0 +1,20 @@
+"""Defines a common interface for the K-Scale Store API."""
+
+from pathlib import Path
+
+from kscale.store.urdf import download_urdf
+from kscale.utils.api_base import APIBase
+
+
+class StoreAPI(APIBase):
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+    ) -> None:
+        super().__init__()
+
+        self.api_key = api_key
+
+    async def urdf(self, artifact_id: str) -> Path:
+        return await download_urdf(artifact_id)

kscale/store/cli.py

@@ -1,12 +1,13 @@
 """Defines the top-level KOL CLI."""
 
 import argparse
+import asyncio
 from typing import Sequence
 
 from kscale.store import pybullet, urdf
 
 
-def main(args: Sequence[str] | None = None) -> None:
+async def main(args: Sequence[str] | None = None) -> None:
     parser = argparse.ArgumentParser(description="K-Scale OnShape Library", add_help=False)
     parser.add_argument(
         "subcommand",
@@ -20,11 +21,15 @@ def main(args: Sequence[str] | None = None) -> None:
 
     match parsed_args.subcommand:
         case "urdf":
-            urdf.main(remaining_args)
+            await urdf.main(remaining_args)
         case "pybullet":
-            pybullet.main(remaining_args)
+            await pybullet.main(remaining_args)
+
+
+def sync_main(args: Sequence[str] | None = None) -> None:
+    asyncio.run(main(args))
 
 
 if __name__ == "__main__":
     # python3 -m kscale.store.cli
-    main()
+    sync_main()

kscale/store/client.py

@@ -0,0 +1,79 @@
+"""Defines a typed client for the K-Scale Store API."""
+
+import logging
+from types import TracebackType
+from typing import Any, Dict, Type
+from urllib.parse import urljoin
+
+import httpx
+from pydantic import BaseModel
+
+from kscale.store.gen.api import (
+    NewListingRequest,
+    NewListingResponse,
+    SingleArtifactResponse,
+    UploadArtifactResponse,
+)
+from kscale.store.utils import API_ROOT, get_api_key
+
+logger = logging.getLogger(__name__)
+
+
+class KScaleStoreClient:
+    def __init__(self, base_url: str = API_ROOT) -> None:
+        self.base_url = base_url
+        self.client = httpx.AsyncClient(
+            base_url=self.base_url,
+            headers={"Authorization": f"Bearer {get_api_key()}"},
+        )
+
+    async def _request(
+        self,
+        method: str,
+        endpoint: str,
+        *,
+        params: Dict[str, Any] | None = None,
+        data: BaseModel | None = None,
+        files: Dict[str, Any] | None = None,
+    ) -> Dict[str, Any]:
+        url = urljoin(self.base_url, endpoint)
+        kwargs: Dict[str, Any] = {"params": params}
+
+        if data:
+            kwargs["json"] = data.dict(exclude_unset=True)
+        if files:
+            kwargs["files"] = files
+
+        response = await self.client.request(method, url, **kwargs)
+        if response.is_error:
+            logger.error(f"Error response from K-Scale Store: {response.text}")
+        response.raise_for_status()
+        return response.json()
+
+    async def get_artifact_info(self, artifact_id: str) -> SingleArtifactResponse:
+        data = await self._request("GET", f"/artifacts/info/{artifact_id}")
+        return SingleArtifactResponse(**data)
+
+    async def upload_artifact(self, listing_id: str, file_path: str) -> UploadArtifactResponse:
+        with open(file_path, "rb") as f:
+            files = {"files": (f.name, f, "application/gzip")}
+            data = await self._request("POST", f"/artifacts/upload/{listing_id}", files=files)
+        return UploadArtifactResponse(**data)
+
+    async def create_listing(self, request: NewListingRequest) -> NewListingResponse:
+        data = await self._request("POST", "/listings", data=request)
+        return NewListingResponse(**data)
+
+    async def close(self) -> None:
+        await self.client.aclose()
+
+    async def __aenter__(self) -> "KScaleStoreClient":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
+        await self.close()

kscale/store/gen/api.py

@@ -2,12 +2,12 @@
 
 # generated by datamodel-codegen:
 #   filename:  openapi.json
-#   timestamp: 2024-08-19T06:07:36+00:00
+#   timestamp: 2024-09-04T04:33:58+00:00
 
 from __future__ import annotations
 
 from enum import Enum
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union
 
 from pydantic import BaseModel, EmailStr, Field
 
@@ -21,8 +21,9 @@ class AuthResponse(BaseModel):
     api_key: str = Field(..., title="Api Key")
 
 
-class BodySetUrdfUrdfUploadListingIdPost(BaseModel):
-    file: bytes = Field(..., title="File")
+class BodyPullOnshapeDocumentOnshapePullListingIdGet(BaseModel):
+    suffix_to_joint_effort: Optional[Dict[str, float]] = Field(None, title="Suffix To Joint Effort")
+    suffix_to_joint_velocity: Optional[Dict[str, float]] = Field(None, title="Suffix To Joint Velocity")
 
 
 class BodyUploadArtifactsUploadListingIdPost(BaseModel):
@@ -51,7 +52,13 @@ class GetListingResponse(BaseModel):
     description: Optional[str] = Field(..., title="Description")
     child_ids: List[str] = Field(..., title="Child Ids")
     tags: List[str] = Field(..., title="Tags")
+    onshape_url: Optional[str] = Field(..., title="Onshape Url")
     can_edit: bool = Field(..., title="Can Edit")
+    created_at: int = Field(..., title="Created At")
+    views: int = Field(..., title="Views")
+    score: int = Field(..., title="Score")
+    user_vote: Optional[bool] = Field(..., title="User Vote")
+    creator_id: str = Field(..., title="Creator Id")
 
 
 class GetTokenResponse(BaseModel):
@@ -82,42 +89,6 @@ class KeysResponseItem(BaseModel):
     permissions: Optional[List[Permission]] = Field(..., title="Permissions")
 
 
-class ArtifactType(Enum):
-    image = "image"
-
-
-class ArtifactType1(Enum):
-    urdf = "urdf"
-    mjcf = "mjcf"
-
-
-class ArtifactType2(Enum):
-    stl = "stl"
-    obj = "obj"
-    dae = "dae"
-    ply = "ply"
-
-
-class ArtifactType3(Enum):
-    tgz = "tgz"
-    zip = "zip"
-
-
-class ListArtifactsItem(BaseModel):
-    artifact_id: str = Field(..., title="Artifact Id")
-    listing_id: str = Field(..., title="Listing Id")
-    name: str = Field(..., title="Name")
-    artifact_type: Union[ArtifactType, ArtifactType1, ArtifactType2, ArtifactType3] = Field(..., title="Artifact Type")
-    description: Optional[str] = Field(..., title="Description")
-    timestamp: int = Field(..., title="Timestamp")
-    urls: ArtifactUrls
-    is_new: Optional[bool] = Field(None, title="Is New")
-
-
-class ListArtifactsResponse(BaseModel):
-    artifacts: List[ListArtifactsItem] = Field(..., title="Artifacts")
-
-
 class ListKeysResponse(BaseModel):
     keys: List[KeysResponseItem] = Field(..., title="Keys")
 
@@ -130,9 +101,16 @@ class ListListingsResponse(BaseModel):
 class Listing(BaseModel):
     id: str = Field(..., title="Id")
     user_id: str = Field(..., title="User Id")
+    created_at: int = Field(..., title="Created At")
+    updated_at: int = Field(..., title="Updated At")
     name: str = Field(..., title="Name")
     child_ids: List[str] = Field(..., title="Child Ids")
-    description: Optional[str] = Field(..., title="Description")
+    description: Optional[str] = Field(None, title="Description")
+    onshape_url: Optional[str] = Field(None, title="Onshape Url")
+    views: Optional[int] = Field(0, title="Views")
+    upvotes: Optional[int] = Field(0, title="Upvotes")
+    downvotes: Optional[int] = Field(0, title="Downvotes")
+    score: Optional[int] = Field(0, title="Score")
 
 
 class ListingInfoResponse(BaseModel):
@@ -141,6 +119,11 @@ class ListingInfoResponse(BaseModel):
     description: Optional[str] = Field(..., title="Description")
     child_ids: List[str] = Field(..., title="Child Ids")
     image_url: Optional[str] = Field(..., title="Image Url")
+    onshape_url: Optional[str] = Field(..., title="Onshape Url")
+    created_at: int = Field(..., title="Created At")
+    views: int = Field(..., title="Views")
+    score: int = Field(..., title="Score")
+    user_vote: Optional[bool] = Field(..., title="User Vote")
 
 
 class LoginRequest(BaseModel):
@@ -163,6 +146,10 @@ class MyUserInfoResponse(BaseModel):
     github_id: Optional[str] = Field(..., title="Github Id")
     google_id: Optional[str] = Field(..., title="Google Id")
     permissions: Optional[List[Permission1]] = Field(..., title="Permissions")
+    first_name: Optional[str] = Field(..., title="First Name")
+    last_name: Optional[str] = Field(..., title="Last Name")
+    name: Optional[str] = Field(..., title="Name")
+    bio: Optional[str] = Field(..., title="Bio")
 
 
 class NewKeyRequest(BaseModel):
@@ -200,6 +187,48 @@ class PublicUsersInfoResponse(BaseModel):
     users: List[PublicUserInfoResponseItem] = Field(..., title="Users")
 
 
+class SetRequest(BaseModel):
+    onshape_url: Optional[str] = Field(..., title="Onshape Url")
+
+
+class ArtifactType(Enum):
+    image = "image"
+
+
+class ArtifactType1(Enum):
+    urdf = "urdf"
+    mjcf = "mjcf"
+
+
+class ArtifactType2(Enum):
+    stl = "stl"
+    obj = "obj"
+    dae = "dae"
+    ply = "ply"
+
+
+class ArtifactType3(Enum):
+    tgz = "tgz"
+    zip = "zip"
+
+
+class SingleArtifactResponse(BaseModel):
+    artifact_id: str = Field(..., title="Artifact Id")
+    listing_id: str = Field(..., title="Listing Id")
+    name: str = Field(..., title="Name")
+    artifact_type: Union[ArtifactType, ArtifactType1, ArtifactType2, ArtifactType3] = Field(..., title="Artifact Type")
+    description: Optional[str] = Field(..., title="Description")
+    timestamp: int = Field(..., title="Timestamp")
+    urls: ArtifactUrls
+    is_new: Optional[bool] = Field(None, title="Is New")
+
+
+class SortOption(Enum):
+    newest = "newest"
+    most_viewed = "most_viewed"
+    most_upvoted = "most_upvoted"
+
+
 class UpdateArtifactRequest(BaseModel):
     name: Optional[str] = Field(None, title="Name")
     description: Optional[str] = Field(None, title="Description")
@@ -212,18 +241,19 @@ class UpdateListingRequest(BaseModel):
     tags: Optional[List[str]] = Field(None, title="Tags")
 
 
-class UploadArtifactResponse(BaseModel):
-    artifacts: List[ListArtifactsItem] = Field(..., title="Artifacts")
-
-
-class UrdfInfo(BaseModel):
-    artifact_id: str = Field(..., title="Artifact Id")
-    url: str = Field(..., title="Url")
+class UpdateUserRequest(BaseModel):
+    email: Optional[str] = Field(None, title="Email")
+    password: Optional[str] = Field(None, title="Password")
+    github_id: Optional[str] = Field(None, title="Github Id")
+    google_id: Optional[str] = Field(None, title="Google Id")
+    first_name: Optional[str] = Field(None, title="First Name")
+    last_name: Optional[str] = Field(None, title="Last Name")
+    name: Optional[str] = Field(None, title="Name")
+    bio: Optional[str] = Field(None, title="Bio")
 
 
-class UrdfResponse(BaseModel):
-    urdf: Optional[UrdfInfo]
-    listing_id: str = Field(..., title="Listing Id")
+class UploadArtifactResponse(BaseModel):
+    artifacts: List[SingleArtifactResponse] = Field(..., title="Artifacts")
 
 
 class UserInfoResponseItem(BaseModel):
@@ -235,7 +265,7 @@ class UserPublic(BaseModel):
     id: str = Field(..., title="Id")
     email: str = Field(..., title="Email")
     permissions: Optional[List[Permission1]] = Field(None, title="Permissions")
-    created_at: Optional[int] = Field(None, title="Created At")
+    created_at: int = Field(..., title="Created At")
     updated_at: Optional[int] = Field(None, title="Updated At")
     first_name: Optional[str] = Field(None, title="First Name")
     last_name: Optional[str] = Field(None, title="Last Name")
@@ -265,3 +295,7 @@ class GetBatchListingsResponse(BaseModel):
 
 class HTTPValidationError(BaseModel):
     detail: Optional[List[ValidationError]] = Field(None, title="Detail")
+
+
+class ListArtifactsResponse(BaseModel):
+    artifacts: List[SingleArtifactResponse] = Field(..., title="Artifacts")

kscale/store/pybullet.py

@@ -1,6 +1,7 @@
 """Simple script to interact with a URDF in PyBullet."""
 
 import argparse
+import asyncio
 import itertools
 import logging
 import math
@@ -8,12 +9,14 @@ import time
 from pathlib import Path
 from typing import Sequence
 
+from kscale.store.urdf import download_urdf
+
 logger = logging.getLogger(__name__)
 
 
-def main(args: Sequence[str] | None = None) -> None:
+async def main(args: Sequence[str] | None = None) -> None:
     parser = argparse.ArgumentParser(description="Show a URDF in PyBullet")
-    parser.add_argument("urdf", nargs="?", help="Path to the URDF file")
+    parser.add_argument("listing_id", help="Listing ID for the URDF")
     parser.add_argument("--dt", type=float, default=0.01, help="Time step")
     parser.add_argument("-n", "--hide-gui", action="store_true", help="Hide the GUI")
     parser.add_argument("--no-merge", action="store_true", help="Do not merge fixed links")
@@ -23,6 +26,9 @@ def main(args: Sequence[str] | None = None) -> None:
     parser.add_argument("--show-collision", action="store_true", help="Show collision meshes")
     parsed_args = parser.parse_args(args)
 
+    # Gets the URDF path.
+    urdf_path = await download_urdf(parsed_args.listing_id)
+
     try:
         import pybullet as p  # type: ignore[import-not-found]
     except ImportError:
@@ -46,13 +52,6 @@ def main(args: Sequence[str] | None = None) -> None:
     # Loads the floor plane.
     floor = p.loadURDF(str((Path(__file__).parent / "bullet" / "plane.urdf").resolve()))
 
-    urdf_path = Path("robot" if parsed_args.urdf is None else parsed_args.urdf)
-    if urdf_path.is_dir():
-        try:
-            urdf_path = next(urdf_path.glob("*.urdf"))
-        except StopIteration:
-            raise FileNotFoundError(f"No URDF files found in {urdf_path}")
-
     # Load the robot URDF.
     start_position = [0.0, 0.0, 1.0]
     start_orientation = p.getQuaternionFromEuler([0.0, 0.0, 0.0])
@@ -175,4 +174,4 @@ def main(args: Sequence[str] | None = None) -> None:
 
 if __name__ == "__main__":
     # python -m kscale.store.pybullet
-    main()
+    asyncio.run(main())

kscale/store/urdf.py

@@ -3,167 +3,169 @@
 import argparse
 import asyncio
 import logging
-import os
+import shutil
 import sys
 import tarfile
 from pathlib import Path
-from typing import Literal, Sequence
+from typing import Literal, Sequence, get_args
 
 import httpx
 import requests
 
 from kscale.conf import Settings
-from kscale.store.gen.api import UrdfResponse
+from kscale.store.client import KScaleStoreClient
+from kscale.store.gen.api import SingleArtifactResponse
+from kscale.store.utils import get_api_key
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 
-def get_api_key() -> str:
-    api_key = Settings.load().store.api_key
-    if not api_key:
-        raise ValueError(
-            "API key not found! Get one here and set it as the `KSCALE_API_KEY` environment variable or in your"
-            "config file: https://kscale.store/keys"
-        )
-    return api_key
-
-
 def get_cache_dir() -> Path:
     return Path(Settings.load().store.cache_dir).expanduser().resolve()
 
 
-def fetch_urdf_info(listing_id: str) -> UrdfResponse:
-    url = f"https://api.kscale.store/urdf/info/{listing_id}"
-    headers = {
-        "Authorization": f"Bearer {get_api_key()}",
-    }
-    response = requests.get(url, headers=headers)
-    response.raise_for_status()
-    return UrdfResponse(**response.json())
+def get_artifact_dir(artifact_id: str) -> Path:
+    cache_dir = get_cache_dir() / artifact_id
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir
 
 
-async def download_artifact(artifact_url: str, cache_dir: Path) -> str:
-    filename = os.path.join(cache_dir, artifact_url.split("/")[-1])
+async def fetch_urdf_info(artifact_id: str, cache_dir: Path) -> SingleArtifactResponse:
+    response_path = cache_dir / "response.json"
+    if response_path.exists():
+        return SingleArtifactResponse.model_validate_json(response_path.read_text())
+    async with KScaleStoreClient() as client:
+        response = await client.get_artifact_info(artifact_id)
+    response_path.write_text(response.model_dump_json())
+    return response
+
+
+async def download_artifact(artifact_url: str, cache_dir: Path) -> Path:
+    filename = cache_dir / Path(artifact_url).name
     headers = {
         "Authorization": f"Bearer {get_api_key()}",
     }
 
-    if not os.path.exists(filename):
-        logger.info("Downloading artifact from %s" % artifact_url)
+    if not filename.exists():
+        logger.info("Downloading artifact from %s", artifact_url)
 
         async with httpx.AsyncClient() as client:
             response = await client.get(artifact_url, headers=headers)
             response.raise_for_status()
-            with open(filename, "wb") as f:
-                for chunk in response.iter_bytes(chunk_size=8192):
-                    f.write(chunk)
-            logger.info("Artifact downloaded to %s" % filename)
+            filename.write_bytes(response.content)
+            logger.info("Artifact downloaded to %s", filename)
     else:
-        logger.info("Artifact already cached at %s" % filename)
+        logger.info("Artifact already cached at %s", filename)
 
     # Extract the .tgz file
-    extract_dir = os.path.join(cache_dir, os.path.splitext(os.path.basename(filename))[0])
-    if not os.path.exists(extract_dir):
-        logger.info(f"Extracting {filename} to {extract_dir}")
+    extract_dir = cache_dir / filename.stem
+    if not extract_dir.exists():
+        logger.info("Extracting %s to %s", filename, extract_dir)
         with tarfile.open(filename, "r:gz") as tar:
             tar.extractall(path=extract_dir)
-        logger.info("Extraction complete")
     else:
-        logger.info("Artifact already extracted at %s" % extract_dir)
+        logger.info("Artifact already extracted at %s", extract_dir)
 
     return extract_dir
 
 
-def create_tarball(folder_path: str | Path, output_filename: str, cache_dir: Path) -> str:
-    tarball_path = os.path.join(cache_dir, output_filename)
+def create_tarball(folder_path: Path, output_filename: str, cache_dir: Path) -> Path:
+    tarball_path = cache_dir / output_filename
     with tarfile.open(tarball_path, "w:gz") as tar:
-        for root, _, files in os.walk(folder_path):
-            for file in files:
-                file_path = os.path.join(root, file)
-                arcname = os.path.relpath(file_path, start=folder_path)
-                tar.add(file_path, arcname=arcname)
-                logger.info("Added %s as %s" % (file_path, arcname))
-    logger.info("Created tarball %s" % tarball_path)
+        for file_path in folder_path.rglob("*"):
+            if file_path.is_file() and file_path.suffix.lower() in (".urdf", ".mjcf", ".stl", ".obj", ".dae"):
+                tar.add(file_path, arcname=file_path.relative_to(folder_path))
+                logger.info("Added %s to tarball", file_path)
+            else:
+                logger.warning("Skipping %s", file_path)
+    logger.info("Created tarball %s", tarball_path)
     return tarball_path
 
 
-async def upload_artifact(tarball_path: str, listing_id: str) -> None:
-    url = f"https://api.kscale.store/urdf/upload/{listing_id}"
-    headers = {
-        "Authorization": f"Bearer {get_api_key()}",
-    }
+async def download_urdf(artifact_id: str) -> Path:
+    cache_dir = get_artifact_dir(artifact_id)
+    try:
+        urdf_info = await fetch_urdf_info(artifact_id, cache_dir)
+        artifact_url = urdf_info.urls.large
+        return await download_artifact(artifact_url, cache_dir)
+
+    except requests.RequestException:
+        logger.exception("Failed to fetch URDF info")
+        raise
+
+
+async def show_urdf_info(artifact_id: str) -> None:
+    try:
+        urdf_info = await fetch_urdf_info(artifact_id, get_artifact_dir(artifact_id))
+        logger.info("URDF Artifact ID: %s", urdf_info.artifact_id)
+        logger.info("URDF URL: %s", urdf_info.urls.large)
+    except requests.RequestException:
+        logger.exception("Failed to fetch URDF info")
+        raise
+
+
+async def remove_local_urdf(artifact_id: str) -> None:
+    try:
+        if artifact_id.lower() == "all":
+            cache_dir = get_cache_dir()
+            if cache_dir.exists():
+                logger.info("Removing all local caches at %s", cache_dir)
+                shutil.rmtree(cache_dir)
+            else:
+                logger.error("No local caches found")
+        else:
+            artifact_dir = get_artifact_dir(artifact_id)
+            if artifact_dir.exists():
+                logger.info("Removing local cache at %s", artifact_dir)
+                shutil.rmtree(artifact_dir)
+            else:
+                logger.error("No local cache found for artifact %s", artifact_id)
+
+    except Exception:
+        logger.error("Failed to remove local cache")
+        raise
+
+
+async def upload_urdf(listing_id: str, args: Sequence[str]) -> None:
+    parser = argparse.ArgumentParser(description="K-Scale URDF Store", add_help=False)
+    parser.add_argument("root_dir", type=Path, help="The path to the root directory to upload")
+    parsed_args = parser.parse_args(args)
 
-    async with httpx.AsyncClient() as client:
-        with open(tarball_path, "rb") as f:
-            files = {"file": (f.name, f, "application/gzip")}
-            response = await client.post(url, headers=headers, files=files)
+    root_dir = parsed_args.root_dir
+    tarball_path = create_tarball(root_dir, "robot.tgz", get_artifact_dir(listing_id))
+
+    async with KScaleStoreClient() as client:
+        response = await client.upload_artifact(listing_id, str(tarball_path))
+
+    logger.info("Uploaded artifacts: %s", [artifact.artifact_id for artifact in response.artifacts])
 
-            response.raise_for_status()
 
-    logger.info("Uploaded artifact to %s" % url)
+Command = Literal["download", "info", "upload", "remove-local"]
 
 
-def main(args: Sequence[str] | None = None) -> None:
+async def main(args: Sequence[str] | None = None) -> None:
     parser = argparse.ArgumentParser(description="K-Scale URDF Store", add_help=False)
-    parser.add_argument(
-        "command",
-        choices=["get", "info", "upload"],
-        help="The command to run",
-    )
-    parser.add_argument("listing_id", help="The listing ID to operate on")
+    parser.add_argument("command", choices=get_args(Command), help="The command to run")
+    parser.add_argument("id", help="The ID to use (artifact when downloading, listing when uploading)")
     parsed_args, remaining_args = parser.parse_known_args(args)
 
-    command: Literal["get", "info", "upload"] = parsed_args.command
-    listing_id: str = parsed_args.listing_id
-
-    def get_listing_dir() -> Path:
-        (cache_dir := get_cache_dir() / listing_id).mkdir(parents=True, exist_ok=True)
-        return cache_dir
+    command: Command = parsed_args.command
+    id: str = parsed_args.id
 
     match command:
-        case "get":
-            try:
-                urdf_info = fetch_urdf_info(listing_id)
-
-                if urdf_info.urdf:
-                    artifact_url = urdf_info.urdf.url
-                    asyncio.run(download_artifact(artifact_url, get_listing_dir()))
-                else:
-                    logger.info("No URDF found for listing %s" % listing_id)
-            except requests.RequestException as e:
-                logger.error("Failed to fetch URDF info: %s" % e)
-                sys.exit(1)
+        case "download":
+            await download_urdf(id)
 
         case "info":
-            try:
-                urdf_info = fetch_urdf_info(listing_id)
-
-                if urdf_info.urdf:
-                    logger.info("URDF Artifact ID: %s" % urdf_info.urdf.artifact_id)
-                    logger.info("URDF URL: %s" % urdf_info.urdf.url)
-                else:
-                    logger.info("No URDF found for listing %s" % listing_id)
-            except requests.RequestException as e:
-                logger.error("Failed to fetch URDF info: %s" % e)
-                sys.exit(1)
-
-        case "upload":
-            parser = argparse.ArgumentParser(description="Upload a URDF artifact to the K-Scale store")
-            parser.add_argument("folder_path", help="The path to the folder containing the URDF files")
-            parsed_args = parser.parse_args(remaining_args)
-            folder_path = Path(parsed_args.folder_path).expanduser().resolve()
+            await show_urdf_info(id)
 
-            output_filename = f"{listing_id}.tgz"
-            tarball_path = create_tarball(folder_path, output_filename, get_listing_dir())
+        case "remove-local":
+            await remove_local_urdf(id)
 
-            try:
-                urdf_info = fetch_urdf_info(listing_id)
-                asyncio.run(upload_artifact(tarball_path, listing_id))
-            except requests.RequestException as e:
-                logger.error("Failed to upload artifact: %s" % e)
-                sys.exit(1)
+        case "upload":
+            await upload_urdf(id, remaining_args)
 
         case _:
             logger.error("Invalid command")
@@ -171,4 +173,4 @@ def main(args: Sequence[str] | None = None) -> None:
 
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())

kscale/store/utils.py

@@ -0,0 +1,19 @@
+"""Utility functions for interacting with the K-Scale Store API."""
+
+import os
+
+from kscale.conf import Settings
+
+API_ROOT = "https://api.kscale.store"
+
+
+def get_api_key() -> str:
+    api_key = Settings.load().store.api_key
+    if api_key is None:
+        api_key = os.getenv("KSCALE_API_KEY")
+    if not api_key:
+        raise ValueError(
+            "API key not found! Get one here and set it as the `KSCALE_API_KEY` environment variable or in your "
+            "config file: https://kscale.store/keys"
+        )
+    return api_key

kscale/utils/api_base.py

@@ -0,0 +1,6 @@
+"""Defines the base class for the K-Scale API."""
+
+
+class APIBase:
+    def __init__(self) -> None:
+        pass

{kscale-0.0.2.dist-info → kscale-0.0.5.dist-info}/METADATA

@@ -1,15 +1,16 @@
 Metadata-Version: 2.1
 Name: kscale
-Version: 0.0.2
+Version: 0.0.5
 Summary: The kscale project
 Home-page: https://github.com/kscalelabs/kscale
 Author: Benjamin Bolte
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: omegaconf
-Requires-Dist: httpx
-Requires-Dist: requests
+Requires-Dist: omegaconf ==2.3.0
+Requires-Dist: email-validator ==2.2.0
+Requires-Dist: httpx ==0.27.0
+Requires-Dist: requests ==2.32.2
 Provides-Extra: dev
 Requires-Dist: black ; extra == 'dev'
 Requires-Dist: darglint ; extra == 'dev'
@@ -50,3 +51,32 @@ This is a command line tool for interacting with various services provided by K-
 ```bash
 pip install kscale
 ```
+
+## Usage
+
+### CLI
+
+Download a URDF from the K-Scale Store:
+
+```bash
+kscale urdf download <artifact_id>
+```
+
+Upload a URDF to the K-Scale Store:
+
+```bash
+kscale urdf upload <artifact_id> <root_dir>
+```
+
+### Python API
+
+Reference a URDF by ID from the K-Scale Store:
+
+```python
+from kscale import KScale
+
+async def main():
+  kscale = KScale()
+  urdf_dir_path = await kscale.store.urdf("123456")
+  print(urdf_dir_path)
+```

kscale-0.0.5.dist-info/RECORD

@@ -0,0 +1,21 @@
+kscale/__init__.py,sha256=e1NdSlxsBKi5qyOIjXC3vAPQXleZwh2hyVlLHcKQsq8,117
+kscale/api.py,sha256=xBtKj8rgZ400r1Xx9LRY0AzSgIIttoXdejmhHhdVGS0,333
+kscale/conf.py,sha256=9fShFaYTbnrm_eiGjmy8ZtC4Q4m6PQkWPyoF3eNyov8,1424
+kscale/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kscale/store/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kscale/store/api.py,sha256=L-QZMDeOsNo8ovIRjwfBlidYKD510hdoKrmrue7C5Wk,454
+kscale/store/cli.py,sha256=8ygg_1tZzOOHJotEIgSN9pfumcriPmA31sI_FCFQiTo,859
+kscale/store/client.py,sha256=HEyTzgCOXdLCQKdFKwNglvTBtXVipul7cxbdjz53Pb4,2544
+kscale/store/pybullet.py,sha256=viTQCE2jT72miPKZpFKj4f4zGLLYtjhRFxGqzxd2Z8M,7504
+kscale/store/urdf.py,sha256=zMUfJNNpZkoLf-SnhMO28f8Qj39hmZrA0vqLL3sZSqE,6069
+kscale/store/utils.py,sha256=vQSFd9fByDUUSD0dAA65T_WI5R55vY-HOQjIJg_u2jw,536
+kscale/store/gen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kscale/store/gen/api.py,sha256=82D41J6pg9KWdgD0lx7NggLcNS32SpnN8DqE3Md6ON0,9559
+kscale/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+kscale/utils/api_base.py,sha256=Kk_WtRDdJHmOg6NtHmVxVrcfARSUkhfr29ypLch_pO0,112
+kscale-0.0.5.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
+kscale-0.0.5.dist-info/METADATA,sha256=jAkpxXCrUa8b30BY32OaW6eIoR0VFHQzKgzO26EDFGM,2514
+kscale-0.0.5.dist-info/WHEEL,sha256=uCRv0ZEik_232NlR4YDw4Pv3Ajt5bKvMH13NUU7hFuI,91
+kscale-0.0.5.dist-info/entry_points.txt,sha256=PaVs1ivqB0BBdGUsiFkxGUYjGLz05VqagxwRVwi4yV4,54
+kscale-0.0.5.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
+kscale-0.0.5.dist-info/RECORD,,

{kscale-0.0.2.dist-info → kscale-0.0.5.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (73.0.1)
+Generator: setuptools (74.1.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

kscale-0.0.5.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+kscale = kscale.store.cli:sync_main

kscale-0.0.2.dist-info/RECORD

@@ -1,15 +0,0 @@
-kscale/__init__.py,sha256=QvlVh4JTl3JL7jQAja76yKtT-IvF4631ASjWY1wS6AQ,22
-kscale/conf.py,sha256=dnO8qii7JeMPMdjfuxnFXURRM4krwzL404RnwDkyL2M,1441
-kscale/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kscale/store/__init__.py,sha256=sXLh7g3KC4QCFxcZGBTpG2scR7hmmBsMjq6LqRptkRg,22
-kscale/store/cli.py,sha256=sf6F8ZLMenNpxi9dPB6kpaUOmRdfNBcKSnAcoN7EMo0,733
-kscale/store/pybullet.py,sha256=2Pog9wPSyyhmhTJY6x5KhuUQgenpTPIvp_9wxOy9ZaU,7622
-kscale/store/urdf.py,sha256=-dp0DEMoDNIqieG9egPv_XaV7Cg0GthMI_QIYvG2Xjk,6185
-kscale/store/gen/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kscale/store/gen/api.py,sha256=7r6KvlZyB-kWKSCixc-YOtaWPy2mn0v11d2tuLlRMSc,7700
-kscale-0.0.2.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-kscale-0.0.2.dist-info/METADATA,sha256=fNvMO1rM9-UfnHjp3BZ4t4FmcQ-CRPAmBDAdXRkBu-8,2028
-kscale-0.0.2.dist-info/WHEEL,sha256=Mdi9PDNwEZptOjTlUcAth7XJDFtKrHYaQMPulZeBCiQ,91
-kscale-0.0.2.dist-info/entry_points.txt,sha256=Vfj9O643497OONpgwy5UJLJ5a5Q1CfHZSuYKDB9D_GI,49
-kscale-0.0.2.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
-kscale-0.0.2.dist-info/RECORD,,

kscale-0.0.2.dist-info/entry_points.txt

@@ -1,2 +0,0 @@
-[console_scripts]
-kscale = kscale.store.cli:main