kscale 0.0.13__cp311-cp311-macosx_11_0_arm64.whl

kscale/__init__.py

@@ -0,0 +1,9 @@
+"""Defines the common interface for the K-Scale Python API."""
+
+__version__ = "0.0.13"
+
+from pathlib import Path
+
+from kscale.api import K
+
+ROOT_DIR = Path(__file__).parent

kscale/api.py

@@ -0,0 +1,14 @@
+"""Defines common functionality for the K-Scale API."""
+
+from kscale.utils.api_base import APIBase
+from kscale.web.api import WebAPI
+
+
+class K(
+    WebAPI,
+    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/artifacts/__init__.py

@@ -0,0 +1,8 @@
+"""Defines some helper functions for working with artifacts."""
+
+from pathlib import Path
+
+ARTIFACTS_DIR = Path(__file__).parent.resolve()
+
+PLANE_OBJ_PATH = ARTIFACTS_DIR / "plane.obj"
+PLANE_URDF_PATH = ARTIFACTS_DIR / "plane.urdf"

kscale/artifacts/plane.obj

@@ -0,0 +1,18 @@
+# Blender v2.66 (sub 1) OBJ File: ''
+# www.blender.org
+mtllib plane.mtl
+o Plane
+v 15.000000 -15.000000 0.000000
+v 15.000000 15.000000 0.000000
+v -15.000000 15.000000 0.000000
+v -15.000000  -15.000000 0.000000
+
+vt 15.000000 0.000000
+vt 15.000000 15.000000
+vt 0.000000 15.000000
+vt 0.000000 0.000000
+
+usemtl Material
+s off
+f 1/1 2/2 3/3
+f 1/1 3/3 4/4

kscale/artifacts/plane.urdf

@@ -0,0 +1,28 @@
+<?xml version="0.0" ?>
+<robot name="plane">
+    <link name="planeLink">
+        <contact>
+            <lateral_friction value="1"/>
+        </contact>
+        <inertial>
+            <origin rpy="0 0 0" xyz="0 0 0"/>
+            <mass value=".0"/>
+            <inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0"/>
+        </inertial>
+        <visual>
+            <origin rpy="0 0 0" xyz="0 0 0"/>
+            <geometry>
+                <mesh filename="plane.obj" scale="1 1 1"/>
+            </geometry>
+            <material name="white">
+                <color rgba="1 1 1 1"/>
+            </material>
+        </visual>
+        <collision>
+            <origin rpy="0 0 0" xyz="0 0 -5"/>
+            <geometry>
+                <box size="30 30 10"/>
+            </geometry>
+        </collision>
+    </link>
+</robot>

kscale/cli.py

@@ -0,0 +1,25 @@
+"""Defines the top-level KOL CLI."""
+
+import click
+
+from kscale.utils.cli import recursive_help
+from kscale.web.kernels import cli as kernel_images_cli
+from kscale.web.krec import cli as krec_cli
+from kscale.web.pybullet import cli as pybullet_cli
+from kscale.web.urdf import cli as urdf_cli
+
+
+@click.group()
+def cli() -> None:
+    """Command line interface for interacting with the K-Scale web API."""
+    pass
+
+
+cli.add_command(urdf_cli, "urdf")
+cli.add_command(pybullet_cli, "pybullet")
+cli.add_command(kernel_images_cli, "kernel")
+cli.add_command(krec_cli, "krec")
+
+if __name__ == "__main__":
+    # python -m kscale.cli
+    print(recursive_help(cli))

kscale/conf.py

@@ -0,0 +1,45 @@
+"""Defines the bot environment settings."""
+
+import functools
+import os
+import warnings
+from dataclasses import dataclass, field
+from pathlib import Path
+
+from omegaconf import II, OmegaConf
+
+
+def get_path() -> Path:
+    if "KSCALE_CONFIG_DIR" in os.environ:
+        return Path(os.environ["KSCALE_CONFIG_DIR"]).expanduser().resolve()
+    return Path("~/.kscale/").expanduser().resolve()
+
+
+@dataclass
+class WWWSettings:
+    api_key: str | None = field(default=None)
+    cache_dir: str = field(default=II("oc.env:KSCALE_CACHE_DIR,'~/.kscale/cache/'"))
+
+
+@dataclass
+class Settings:
+    www: WWWSettings = field(default_factory=WWWSettings)
+
+    def save(self) -> None:
+        (dir_path := get_path()).mkdir(parents=True, exist_ok=True)
+        with open(dir_path / "settings.yaml", "w") as f:
+            OmegaConf.save(config=self, f=f)
+
+    @functools.lru_cache
+    @staticmethod
+    def load() -> "Settings":
+        config = OmegaConf.structured(Settings)
+        if not (dir_path := get_path()).exists():
+            warnings.warn(f"Settings directory does not exist: {dir_path}. Creating it now.")
+            dir_path.mkdir(parents=True)
+            OmegaConf.save(config, dir_path / "settings.yaml")
+        else:
+            with open(dir_path / "settings.yaml", "r") as f:
+                raw_settings = OmegaConf.load(f)
+                config = OmegaConf.merge(config, raw_settings)
+        return config

kscale/requirements-dev.txt

@@ -0,0 +1,11 @@
+# requirements-dev.txt
+
+# For linting code.
+black
+darglint
+mypy
+pytest
+ruff
+
+# For generating API types.
+datamodel-code-generator

kscale/requirements.txt

@@ -0,0 +1,17 @@
+# requirements.txt
+
+# Configuration
+omegaconf
+email_validator
+
+# HTTP requests
+httpx
+requests
+pydantic
+
+# CLI
+click
+aiofiles
+
+# K-Scale
+krec

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/utils/checksum.py

@@ -0,0 +1,41 @@
+"""Utility functions for file checksums."""
+
+import hashlib
+from pathlib import Path
+from typing import Tuple
+
+CHUNK_SIZE = 8192
+
+
+async def calculate_sha256(file_path: str | Path) -> Tuple[str, int]:
+    """Calculate SHA256 checksum and size of a file.
+
+    Args:
+        file_path: Path to the file
+
+    Returns:
+        Tuple of (checksum hex string, file size in bytes)
+    """
+    sha256_hash = hashlib.sha256()
+    file_size = 0
+
+    with open(file_path, "rb") as f:
+        for chunk in iter(lambda: f.read(CHUNK_SIZE), b""):
+            sha256_hash.update(chunk)
+            file_size += len(chunk)
+
+    return sha256_hash.hexdigest(), file_size
+
+
+class FileChecksum:
+    """Helper class for handling file checksums."""
+
+    @staticmethod
+    async def calculate(file_path: str | Path) -> Tuple[str, int]:
+        """Calculate SHA256 checksum and size of a file."""
+        return await calculate_sha256(file_path)
+
+    @staticmethod
+    def update_hash(hash_obj: "hashlib._Hash", chunk: bytes) -> None:
+        """Update a hash object with new data."""
+        hash_obj.update(chunk)

kscale/utils/cli.py

@@ -0,0 +1,28 @@
+"""Defines utilities for working with asyncio."""
+
+import asyncio
+import textwrap
+from functools import wraps
+from typing import Any, Callable, Coroutine, ParamSpec, TypeVar
+
+import click
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+
+def coro(f: Callable[P, Coroutine[Any, Any, T]]) -> Callable[P, T]:
+    @wraps(f)
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:
+        return asyncio.run(f(*args, **kwargs))
+
+    return wrapper
+
+
+def recursive_help(cmd: click.Command, parent: click.Context | None = None, indent: int = 0) -> str:
+    ctx = click.core.Context(cmd, info_name=cmd.name, parent=parent)
+    help_text = cmd.get_help(ctx)
+    commands = getattr(cmd, "commands", {})
+    for sub in commands.values():
+        help_text += recursive_help(sub, ctx, indent + 2)
+    return textwrap.indent(help_text, " " * indent)

kscale/web/api.py

@@ -0,0 +1,98 @@
+"""Defines a common interface for the K-Scale WWW API."""
+
+import asyncio
+from pathlib import Path
+from typing import overload
+
+from kscale.utils.api_base import APIBase
+from kscale.web.gen.api import UploadArtifactResponse
+from kscale.web.urdf import download_urdf, upload_urdf
+
+
+class WebAPI(APIBase):
+    def __init__(
+        self,
+        *,
+        api_key: str | None = None,
+    ) -> None:
+        super().__init__()
+
+        self.api_key = api_key
+
+    async def artifact_root(self, artifact_id: str) -> Path:
+        return await download_urdf(artifact_id)
+
+    @overload
+    async def urdf_path(self, artifact_id: str) -> Path: ...
+
+    @overload
+    async def urdf_path(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None: ...
+
+    async def urdf_path(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None:
+        root_dir = await self.artifact_root(artifact_id)
+        urdf_path = next(root_dir.glob("*.urdf"), None)
+        if urdf_path is None and throw_if_missing:
+            raise FileNotFoundError(f"No URDF found for artifact {artifact_id}")
+        return urdf_path
+
+    @overload
+    def urdf_path_sync(self, artifact_id: str) -> Path: ...
+
+    @overload
+    def urdf_path_sync(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None: ...
+
+    def urdf_path_sync(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None:
+        return asyncio.run(self.urdf_path(artifact_id, throw_if_missing=throw_if_missing))
+
+    @overload
+    async def mjcf_path(self, artifact_id: str) -> Path: ...
+
+    @overload
+    async def mjcf_path(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None: ...
+
+    async def mjcf_path(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None:
+        root_dir = await self.artifact_root(artifact_id)
+        mjcf_path = next(root_dir.glob("*.mjcf"), None)
+        if mjcf_path is None and throw_if_missing:
+            raise FileNotFoundError(f"No MJCF found for artifact {artifact_id}")
+        return mjcf_path
+
+    @overload
+    def mjcf_path_sync(self, artifact_id: str) -> Path: ...
+
+    @overload
+    def mjcf_path_sync(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None: ...
+
+    def mjcf_path_sync(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None:
+        return asyncio.run(self.mjcf_path(artifact_id, throw_if_missing=throw_if_missing))
+
+    @overload
+    async def xml_path(self, artifact_id: str) -> Path: ...
+
+    @overload
+    async def xml_path(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None: ...
+
+    async def xml_path(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None:
+        root_dir = await self.artifact_root(artifact_id)
+        xml_path = next(root_dir.glob("*.xml"), None)
+        if xml_path is None and throw_if_missing:
+            raise FileNotFoundError(f"No XML found for artifact {artifact_id}")
+        return xml_path
+
+    async def upload_urdf(self, listing_id: str, root_dir: Path) -> UploadArtifactResponse:
+        return await upload_urdf(listing_id, root_dir)
+
+    def artifact_root_sync(self, artifact_id: str) -> Path:
+        return asyncio.run(self.artifact_root(artifact_id))
+
+    @overload
+    def xml_path_sync(self, artifact_id: str) -> Path: ...
+
+    @overload
+    def xml_path_sync(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None: ...
+
+    def xml_path_sync(self, artifact_id: str, *, throw_if_missing: bool = True) -> Path | None:
+        return asyncio.run(self.xml_path(artifact_id, throw_if_missing=throw_if_missing))
+
+    def upload_urdf_sync(self, listing_id: str, root_dir: Path) -> UploadArtifactResponse:
+        return asyncio.run(self.upload_urdf(listing_id, root_dir))