kscale 0.0.10__py3-none-any.whl → 0.1.0__py3-none-any.whl

kscale/store/pybullet.py

@@ -1,180 +0,0 @@
-"""Simple script to interact with a URDF in PyBullet."""
-
-import argparse
-import asyncio
-import itertools
-import logging
-import math
-import time
-from typing import Sequence
-
-from kscale.artifacts import PLANE_URDF_PATH
-from kscale.store.urdf import download_urdf
-
-logger = logging.getLogger(__name__)
-
-
-async def main(args: Sequence[str] | None = None) -> None:
-    parser = argparse.ArgumentParser(description="Show a URDF in PyBullet")
-    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")
-    parser.add_argument("--hide-origin", action="store_true", help="Do not show the origin")
-    parser.add_argument("--show-inertia", action="store_true", help="Visualizes the inertia frames")
-    parser.add_argument("--see-thru", action="store_true", help="Use see-through mode")
-    parser.add_argument("--show-collision", action="store_true", help="Show collision meshes")
-    parsed_args = parser.parse_args(args)
-
-    # Gets the URDF path.
-    urdf_dir = await download_urdf(parsed_args.listing_id)
-    urdf_path = next(urdf_dir.glob("*.urdf"), None)
-    if urdf_path is None:
-        raise ValueError(f"No URDF found in {urdf_dir}")
-
-    try:
-        import pybullet as p  # type: ignore[import-not-found]
-    except ImportError:
-        raise ImportError("pybullet is required to run this script")
-
-    # Connect to PyBullet.
-    p.connect(p.GUI)
-    p.setGravity(0, 0, -9.81)
-    p.setRealTimeSimulation(0)
-
-    # Turn off panels.
-    if parsed_args.hide_gui:
-        p.configureDebugVisualizer(p.COV_ENABLE_GUI, 0)
-    p.configureDebugVisualizer(p.COV_ENABLE_SEGMENTATION_MARK_PREVIEW, 0)
-    p.configureDebugVisualizer(p.COV_ENABLE_DEPTH_BUFFER_PREVIEW, 0)
-    p.configureDebugVisualizer(p.COV_ENABLE_RGB_BUFFER_PREVIEW, 0)
-
-    # Enable mouse picking.
-    p.configureDebugVisualizer(p.COV_ENABLE_MOUSE_PICKING, 1)
-
-    # Loads the floor plane.
-    floor = p.loadURDF(str(PLANE_URDF_PATH))
-
-    # Load the robot URDF.
-    start_position = [0.0, 0.0, 1.0]
-    start_orientation = p.getQuaternionFromEuler([0.0, 0.0, 0.0])
-    flags = p.URDF_USE_INERTIA_FROM_FILE
-    if not parsed_args.no_merge:
-        flags |= p.URDF_MERGE_FIXED_LINKS
-    robot = p.loadURDF(str(urdf_path), start_position, start_orientation, flags=flags, useFixedBase=0)
-
-    # Display collision meshes as separate object.
-    if parsed_args.show_collision:
-        collision_flags = p.URDF_USE_INERTIA_FROM_FILE | p.URDF_USE_SELF_COLLISION_EXCLUDE_ALL_PARENTS
-        collision = p.loadURDF(str(urdf_path), start_position, start_orientation, flags=collision_flags, useFixedBase=0)
-
-        # Make collision shapes semi-transparent.
-        joint_ids = [i for i in range(p.getNumJoints(collision))] + [-1]
-        for i in joint_ids:
-            p.changeVisualShape(collision, i, rgbaColor=[1, 0, 0, 0.5])
-
-    # Initializes physics parameters.
-    p.changeDynamics(floor, -1, lateralFriction=1, spinningFriction=-1, rollingFriction=-1)
-    p.setPhysicsEngineParameter(fixedTimeStep=parsed_args.dt, maxNumCmdPer1ms=1000)
-
-    # Shows the origin of the robot.
-    if not parsed_args.hide_origin:
-        p.addUserDebugLine([0, 0, 0], [0.1, 0, 0], [1, 0, 0], parentObjectUniqueId=robot, parentLinkIndex=-1)
-        p.addUserDebugLine([0, 0, 0], [0, 0.1, 0], [0, 1, 0], parentObjectUniqueId=robot, parentLinkIndex=-1)
-        p.addUserDebugLine([0, 0, 0], [0, 0, 0.1], [0, 0, 1], parentObjectUniqueId=robot, parentLinkIndex=-1)
-
-    # Make the robot see-through.
-    joint_ids = [i for i in range(p.getNumJoints(robot))] + [-1]
-    if parsed_args.see_thru:
-        for i in joint_ids:
-            p.changeVisualShape(robot, i, rgbaColor=[1, 1, 1, 0.5])
-
-    def draw_box(pt: list[list[float]], color: tuple[float, float, float], obj_id: int, link_id: int) -> None:
-        assert len(pt) == 8
-        assert all(len(p) == 3 for p in pt)
-
-        mapping = [1, 3, 0, 2]
-        for i in range(4):
-            p.addUserDebugLine(pt[i], pt[i + 4], color, 1, parentObjectUniqueId=obj_id, parentLinkIndex=link_id)
-            p.addUserDebugLine(pt[i], pt[mapping[i]], color, 1, parentObjectUniqueId=obj_id, parentLinkIndex=link_id)
-            p.addUserDebugLine(
-                pt[i + 4], pt[mapping[i] + 4], color, 1, parentObjectUniqueId=obj_id, parentLinkIndex=link_id
-            )
-
-    # Shows bounding boxes around each part of the robot representing the inertia frame.
-    if parsed_args.show_inertia:
-        for i in joint_ids:
-            dynamics_info = p.getDynamicsInfo(robot, i)
-            mass = dynamics_info[0]
-            if mass <= 0:
-                continue
-            inertia = dynamics_info[2]
-            ixx = inertia[0]
-            iyy = inertia[1]
-            izz = inertia[2]
-            box_scale_x = 0.5 * math.sqrt(6 * (izz + iyy - ixx) / mass)
-            box_scale_y = 0.5 * math.sqrt(6 * (izz + ixx - iyy) / mass)
-            box_scale_z = 0.5 * math.sqrt(6 * (ixx + iyy - izz) / mass)
-
-            half_extents = [box_scale_x, box_scale_y, box_scale_z]
-            pt = [
-                [x, y, z]
-                for x, y, z in itertools.product(
-                    [-half_extents[0], half_extents[0]],
-                    [-half_extents[1], half_extents[1]],
-                    [-half_extents[2], half_extents[2]],
-                )
-            ]
-            draw_box(pt, (1, 0, 0), robot, i)
-
-    # Show joint controller.
-    joints: dict[str, int] = {}
-    controls: dict[str, float] = {}
-    for i in range(p.getNumJoints(robot)):
-        joint_info = p.getJointInfo(robot, i)
-        name = joint_info[1].decode("utf-8")
-        joint_type = joint_info[2]
-        joints[name] = i
-        if joint_type == p.JOINT_PRISMATIC:
-            joint_min, joint_max = joint_info[8:10]
-            controls[name] = p.addUserDebugParameter(name, joint_min, joint_max, 0.0)
-        elif joint_type == p.JOINT_REVOLUTE:
-            joint_min, joint_max = joint_info[8:10]
-            controls[name] = p.addUserDebugParameter(name, joint_min, joint_max, 0.0)
-
-    # Run the simulation until the user closes the window.
-    last_time = time.time()
-    prev_control_values = {k: 0.0 for k in controls}
-    while p.isConnected():
-        # Reset the simulation if "r" was pressed.
-        keys = p.getKeyboardEvents()
-        if ord("r") in keys and keys[ord("r")] & p.KEY_WAS_TRIGGERED:
-            p.resetBasePositionAndOrientation(robot, start_position, start_orientation)
-            p.setJointMotorControlArray(
-                robot,
-                range(p.getNumJoints(robot)),
-                p.POSITION_CONTROL,
-                targetPositions=[0] * p.getNumJoints(robot),
-            )
-
-        # Set joint positions.
-        for k, v in controls.items():
-            try:
-                target_position = p.readUserDebugParameter(v)
-                if target_position != prev_control_values[k]:
-                    prev_control_values[k] = target_position
-                    p.setJointMotorControl2(robot, joints[k], p.POSITION_CONTROL, target_position)
-            except p.error:
-                logger.debug("Failed to set joint %s", k)
-                pass
-
-        # Step simulation.
-        p.stepSimulation()
-        cur_time = time.time()
-        time.sleep(max(0, parsed_args.dt - (cur_time - last_time)))
-        last_time = cur_time
-
-
-if __name__ == "__main__":
-    # python -m kscale.store.pybullet
-    asyncio.run(main())

kscale/store/urdf.py

@@ -1,193 +0,0 @@
-"""Utility functions for managing artifacts in the K-Scale store."""
-
-import argparse
-import asyncio
-import logging
-import shutil
-import sys
-import tarfile
-from pathlib import Path
-from typing import Literal, Sequence, get_args
-
-import httpx
-import requests
-
-from kscale.conf import Settings
-from kscale.store.client import KScaleStoreClient
-from kscale.store.gen.api import SingleArtifactResponse, UploadArtifactResponse
-from kscale.store.utils import get_api_key
-
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-ALLOWED_SUFFIXES = {
-    ".urdf",
-    ".mjcf",
-    ".stl",
-    ".obj",
-    ".dae",
-    ".png",
-    ".jpg",
-    ".jpeg",
-}
-
-
-def get_cache_dir() -> Path:
-    return Path(Settings.load().store.cache_dir).expanduser().resolve()
-
-
-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 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 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()
-            filename.write_bytes(response.content)
-            logger.info("Artifact downloaded to %s", filename)
-    else:
-        logger.info("Artifact already cached at %s", filename)
-
-    # Extract the .tgz file
-    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)
-    else:
-        logger.info("Artifact already extracted at %s", extract_dir)
-
-    return extract_dir
-
-
-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 file_path in folder_path.rglob("*"):
-            if file_path.is_file() and file_path.suffix.lower() in ALLOWED_SUFFIXES:
-                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 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, root_dir: Path) -> UploadArtifactResponse:
-    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])
-    return response
-
-
-async def upload_urdf_cli(listing_id: str, args: Sequence[str]) -> UploadArtifactResponse:
-    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)
-
-    root_dir = parsed_args.root_dir
-    response = await upload_urdf(listing_id, root_dir)
-    return response
-
-
-Command = Literal["download", "info", "upload", "remove-local"]
-
-
-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_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: Command = parsed_args.command
-    id: str = parsed_args.id
-
-    match command:
-        case "download":
-            await download_urdf(id)
-
-        case "info":
-            await show_urdf_info(id)
-
-        case "remove-local":
-            await remove_local_urdf(id)
-
-        case "upload":
-            await upload_urdf_cli(id, remaining_args)
-
-        case _:
-            logger.error("Invalid command")
-            sys.exit(1)
-
-
-if __name__ == "__main__":
-    asyncio.run(main())

kscale/store/utils.py

@@ -1,33 +0,0 @@
-"""Utility functions for interacting with the K-Scale Store API."""
-
-import os
-
-from kscale.conf import Settings
-
-
-def get_api_root() -> str:
-    """Returns the base URL for the K-Scale Store API.
-
-    This can be overridden when targetting a different server.
-
-    Returns:
-        The base URL for the K-Scale Store API.
-    """
-    return os.getenv("KSCALE_API_ROOT", "https://api.kscale.store")
-
-
-def get_api_key() -> str:
-    """Returns the API key for the K-Scale Store API.
-
-    Returns:
-        The API key for the K-Scale Store API.
-    """
-    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-0.0.10.dist-info/RECORD

@@ -1,26 +0,0 @@
-kscale/__init__.py,sha256=s_FPBlHyXyKecFXm57Q2kH3d0O8suiaE8FrI6n0ckXU,178
-kscale/api.py,sha256=xBtKj8rgZ400r1Xx9LRY0AzSgIIttoXdejmhHhdVGS0,333
-kscale/conf.py,sha256=9fShFaYTbnrm_eiGjmy8ZtC4Q4m6PQkWPyoF3eNyov8,1424
-kscale/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kscale/requirements-dev.txt,sha256=WI7-ea4IRJakmqVMN8QKhOsDGrghwtvk03aIsFaNSIw,130
-kscale/requirements.txt,sha256=gxo_niYIHsmyxKxvIOegv45s_lvdnIzI2iFJ2TZzx_U,103
-kscale/artifacts/__init__.py,sha256=RK8wdybtCJPgdLLJ8R8-YMi1Ph5ojqAKVJZowHONtgo,232
-kscale/artifacts/plane.obj,sha256=x59-IIrWpLjhotChiqT2Ul6U8s0RcHkaEeUZb4KXL1c,348
-kscale/artifacts/plane.urdf,sha256=LCiTk14AyTHjkZ1jvsb0hNaEaJUxDb8Z1JjsgpXu3YM,819
-kscale/store/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kscale/store/api.py,sha256=JBp4h6yz_ESSdB-2FpkD_1-hdI8_iEKh9svzkyy3jhs,2386
-kscale/store/cli.py,sha256=8ygg_1tZzOOHJotEIgSN9pfumcriPmA31sI_FCFQiTo,859
-kscale/store/client.py,sha256=R1IDnf2J4ojAcP8nmUUHfXhcHUt4zP0-mxtVI7MIC5U,2664
-kscale/store/pybullet.py,sha256=zoeATQStuRWgmPhku65xjfgvE3Y8ReheUIAkZnDr2C0,7614
-kscale/store/urdf.py,sha256=5x8tK2BYv901S_yYWYPWEnHv-3T0ALBQMdDwb70EZFw,6395
-kscale/store/utils.py,sha256=rFXGkem2oAttAf3bhWmFEhxrqYnaVvlVJsC268IMw6Y,906
-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.10.dist-info/LICENSE,sha256=HCN2bImAzUOXldAZZI7JZ9PYq6OwMlDAP_PpX1HnuN0,1071
-kscale-0.0.10.dist-info/METADATA,sha256=4K7YAFdI1P5PIAAw78lvbyUw1zIAZ_-ZiZr2NFx-o2Q,2505
-kscale-0.0.10.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-kscale-0.0.10.dist-info/entry_points.txt,sha256=PaVs1ivqB0BBdGUsiFkxGUYjGLz05VqagxwRVwi4yV4,54
-kscale-0.0.10.dist-info/top_level.txt,sha256=C2ynjYwopg6YjgttnI2dJjasyq3EKNmYp-IfQg9Xms4,7
-kscale-0.0.10.dist-info/RECORD,,

kscale-0.0.10.dist-info/entry_points.txt

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