From a6279dee09673d6f449a068223c463c03db37081 Mon Sep 17 00:00:00 2001 From: Benjamin Bolte Date: Sat, 23 Nov 2024 00:22:54 -0800 Subject: [PATCH 1/2] validate krec files --- kscale/requirements.txt | 3 +++ kscale/web/krec.py | 17 ++++++++++++----- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/kscale/requirements.txt b/kscale/requirements.txt index 7da5d5b..6773ecd 100644 --- a/kscale/requirements.txt +++ b/kscale/requirements.txt @@ -12,3 +12,6 @@ pydantic # CLI click aiofiles + +# K-Scale +krec diff --git a/kscale/web/krec.py b/kscale/web/krec.py index c807df4..ce609a0 100644 --- a/kscale/web/krec.py +++ b/kscale/web/krec.py @@ -8,6 +8,7 @@ import aiofiles import click import httpx +import krec from kscale.utils.cli import coro from kscale.web.gen.api import UploadKRecRequest @@ -20,7 +21,6 @@ async def upload_krec( robot_id: str, file_path: Path, - name: str, description: str | None = None, upload_timeout: float = DEFAULT_UPLOAD_TIMEOUT, ) -> str: @@ -32,11 +32,19 @@ async def upload_krec( logger.info("File name: %s", file_path.name) logger.info("File size: %.1f MB", file_size / 1024 / 1024) + if not file_path.suffix.lower() == ".krec": + logger.warning("File extension is not .krec - are you sure this is a valid K-Rec file?") + + try: + krec.KRec.load(file_path) + except Exception as e: + raise ValueError(f"Failed to load K-Rec from {file_path} - are you sure this is a valid K-Rec file?") from e + async with KScaleWWWClient(upload_timeout=upload_timeout) as client: create_response = await client.create_krec( UploadKRecRequest( robot_id=robot_id, - name=name, + name=file_path.name, description=description, ) ) @@ -146,12 +154,11 @@ def cli() -> None: @cli.command() @click.argument("robot_id") @click.argument("file_path", type=click.Path(exists=True, path_type=Path)) -@click.option("--name", "-n", help="Name of the K-Rec", required=True) @click.option("--description", "-d", help="Description of the K-Rec") @coro -async def upload(robot_id: str, file_path: Path, name: str, description: str | None = None) -> None: +async def upload(robot_id: str, file_path: Path, description: str | None = None) -> None: """Upload a K-Rec file.""" - krec_id = await upload_krec(robot_id, file_path, name, description) + krec_id = await upload_krec(robot_id, file_path, description) click.echo(f"Successfully uploaded K-Rec: {krec_id}") From d4e2738f0f19844ce9b1cf40f3ec81532db7580a Mon Sep 17 00:00:00 2001 From: Benjamin Bolte Date: Sat, 23 Nov 2024 00:27:56 -0800 Subject: [PATCH 2/2] fix lint --- kscale/web/krec.py | 6 +++--- pyproject.toml | 1 - 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/kscale/web/krec.py b/kscale/web/krec.py index ce609a0..192327f 100644 --- a/kscale/web/krec.py +++ b/kscale/web/krec.py @@ -36,7 +36,7 @@ async def upload_krec( logger.warning("File extension is not .krec - are you sure this is a valid K-Rec file?") try: - krec.KRec.load(file_path) + krec.KRec.load(str(file_path.resolve())) except Exception as e: raise ValueError(f"Failed to load K-Rec from {file_path} - are you sure this is a valid K-Rec file?") from e @@ -69,8 +69,8 @@ async def upload_krec( return create_response["krec_id"] -def upload_krec_sync(robot_id: str, file_path: Path, name: str, description: str | None = None) -> str: - return asyncio.run(upload_krec(robot_id, file_path, name, description)) +def upload_krec_sync(robot_id: str, file_path: Path, description: str | None = None) -> str: + return asyncio.run(upload_krec(robot_id, file_path, description)) async def fetch_krec_info(krec_id: str, cache_dir: Path) -> dict: diff --git a/pyproject.toml b/pyproject.toml index 1698cd5..fce2114 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -71,7 +71,6 @@ select = [ ] ignore = [ - "ANN101", "ANN102", "D101", "D102", "D103", "D104", "D105", "D106", "D107", "N812", "N817", "PLR0911", "PLR0912", "PLR0913", "PLR0915", "PLR2004",