kleinkram 0.36.3.dev20241113165309__py3-none-any.whl → 0.37.0__py3-none-any.whl

kleinkram/auth.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import urllib.parse
+import webbrowser
+from getpass import getpass
+from http.server import BaseHTTPRequestHandler
+from http.server import HTTPServer
+from typing import Optional
+
+from kleinkram.config import Config
+from kleinkram.config import CONFIG_PATH
+from kleinkram.config import CorruptedConfigFile
+from kleinkram.config import Credentials
+from kleinkram.config import InvalidConfigFile
+
+CLI_CALLBACK_ENDPOINT = "/cli/callback"
+OAUTH_SLUG = "/auth/google?state=cli"
+
+
+def _has_browser() -> bool:
+    try:
+        webbrowser.get()
+        return True
+    except webbrowser.Error:
+        return False
+
+
+def _headless_auth(*, url: str) -> None:
+    config = Config()
+
+    print(f"Please open the following URL manually to authenticate: {url}")
+    print("Enter the authentication token provided after logging in:")
+    auth_token = getpass("Authentication Token: ")
+    refresh_token = getpass("Refresh Token: ")
+
+    if auth_token and refresh_token:
+        creds = Credentials(auth_token=auth_token, refresh_token=refresh_token)
+        config.save_credentials(creds)
+        print(f"Authentication complete. Tokens saved to {CONFIG_PATH}.")
+    else:
+        raise ValueError("Please provided tokens.")
+
+
+class OAuthCallbackHandler(BaseHTTPRequestHandler):
+    def do_GET(self):
+        if self.path.startswith(CLI_CALLBACK_ENDPOINT):
+            query = urllib.parse.urlparse(self.path).query
+            params = urllib.parse.parse_qs(query)
+
+            try:
+                creds = Credentials(
+                    auth_token=params.get("authtoken")[0],  # type: ignore
+                    refresh_token=params.get("refreshtoken")[0],  # type: ignore
+                )
+            except Exception:
+                raise RuntimeError("Failed to fetch authentication tokens.")
+
+            config = Config()
+            config.save_credentials(creds)
+
+            self.send_response(200)
+            self.send_header("Content-type", "text/html")
+            self.end_headers()
+            self.wfile.write(b"Authentication successful. You can close this window.")
+        else:
+            raise RuntimeError("Invalid path")
+
+    def log_message(self, *args, **kwargs):
+        _ = args, kwargs
+        pass  # suppress logging
+
+
+def _browser_auth(*, url: str) -> None:
+    webbrowser.open(url)
+
+    server = HTTPServer(("", 8000), OAuthCallbackHandler)
+    server.handle_request()
+
+    print(f"Authentication complete. Tokens saved to {CONFIG_PATH}.")
+
+
+def login_flow(*, key: Optional[str] = None, headless: bool = False) -> None:
+    config = Config(overwrite=True)
+
+    # use cli key login
+    if key is not None:
+        creds = Credentials(cli_key=key)
+        config.save_credentials(creds)
+
+    url = f"{config.endpoint}{OAUTH_SLUG}"
+
+    if not headless and _has_browser():
+        _browser_auth(url=url)
+    else:
+        headless_url = f"{url}-no-redirect"
+        _headless_auth(url=headless_url)

kleinkram/commands/__init__.py

@@ -0,0 +1 @@
+from __future__ import annotations

kleinkram/commands/download.py

@@ -0,0 +1,103 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from typing import List
+from typing import Optional
+
+import typer
+from kleinkram.api.client import AuthenticatedClient
+from kleinkram.api.file_transfer import download_file
+from kleinkram.api.routes import get_files_by_file_spec
+from kleinkram.config import get_shared_state
+from kleinkram.models import FILE_STATE_COLOR
+from kleinkram.models import files_to_table
+from kleinkram.models import FileState
+from kleinkram.utils import b64_md5
+from kleinkram.utils import get_valid_file_spec
+from kleinkram.utils import to_name_or_uuid
+from rich.console import Console
+
+
+HELP = """\
+Download files from kleinkram.
+"""
+
+
+download_typer = typer.Typer(
+    name="download", no_args_is_help=True, invoke_without_command=True, help=HELP
+)
+
+
+@download_typer.callback()
+def download(
+    files: Optional[List[str]] = typer.Argument(
+        None, help="file names, ids or patterns"
+    ),
+    project: Optional[str] = typer.Option(
+        None, "--project", "-p", help="project name or id"
+    ),
+    mission: Optional[str] = typer.Option(
+        None, "--mission", "-m", help="mission name or id"
+    ),
+    dest: str = typer.Option(prompt="destination", help="local path to save the files"),
+) -> None:
+    _files = [to_name_or_uuid(f) for f in files or []]
+    _project = to_name_or_uuid(project) if project else None
+    _mission = to_name_or_uuid(mission) if mission else None
+
+    # create destionation directory
+    dest_dir = Path(dest)
+
+    if not dest_dir.exists():
+        typer.confirm(f"Destination {dest_dir} does not exist. Create it?", abort=True)
+
+    dest_dir.mkdir(parents=True, exist_ok=True)
+
+    client = AuthenticatedClient()
+    file_spec = get_valid_file_spec(_files, mission=_mission, project=_project)
+    parsed_files = get_files_by_file_spec(client, file_spec)
+
+    # check if filenames are unique
+    if len(set(f.name for f in parsed_files)) != len(parsed_files):
+        raise ValueError(
+            "the files you are trying to download do not have unique names"
+        )
+
+    console = Console()
+    if get_shared_state().verbose:
+        table = files_to_table(parsed_files, title="downloading files...")
+        console.print(table)
+
+    for file in parsed_files:
+        if file.state != FileState.OK:
+            if get_shared_state().verbose:
+                console.print(
+                    f"Skipping file {file.name} with state ",
+                    end="",
+                )
+                console.print(f"{file.state.value}", style=FILE_STATE_COLOR[file.state])
+            else:
+                print(
+                    f"skipping file {file.name} with state {file.state.value}",
+                    file=sys.stderr,
+                )
+            continue
+
+        try:
+            download_file(
+                client,
+                file_id=file.id,
+                name=file.name,
+                dest=dest_dir,
+                hash=file.hash,
+                size=file.size,
+            )
+        except FileExistsError:
+            local_hash = b64_md5(dest_dir / file.name)
+            if local_hash == file.hash:
+                print(f"{file.name} already exists in dest, skipping...")
+            else:
+                print(f"{file.name} already exists in dest, but has different hash!")
+        except Exception as e:
+            print(f"Error downloading file {file.name}: {repr(e)}")

kleinkram/commands/endpoint.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import sys
+
+import typer
+from kleinkram.auth import Config
+
+HELP = """\
+Get or set the current endpoint.
+
+The endpoint is used to determine the API server to connect to\
+(default is the API server of https://datasets.leggedrobotics.com).
+"""
+
+endpoint_typer = typer.Typer(
+    name="endpoint",
+    help=HELP,
+    no_args_is_help=True,
+    context_settings={"help_option_names": ["-h", "--help"]},
+)
+
+
+@endpoint_typer.command("set")
+def set_endpoint(endpoint: str = typer.Argument(None, help="API endpoint to use")):
+    """
+    Use this command to switch between different API endpoints.\n
+    Standard endpoints are:\n
+    - http://localhost:3000\n
+    - https://api.datasets.leggedrobotics.com\n
+    - https://api.datasets.dev.leggedrobotics.com
+    """
+
+    if not endpoint:
+        raise ValueError("No endpoint provided.")
+
+    tokenfile = Config()
+    tokenfile.endpoint = endpoint
+    tokenfile.save()
+
+    print(f"Endpoint set to: {endpoint}")
+    if tokenfile.endpoint not in tokenfile.credentials:
+        print("\nLogin with `klein login`.")
+
+
+@endpoint_typer.command("list")
+def list_endpoints():
+    """
+    Get the current endpoint
+
+    Also displays all endpoints with saved tokens.
+    """
+    config = Config()
+    print(f"Current endpoint: {config.endpoint}\n", file=sys.stderr)
+
+    if not config.credentials:
+        print("No saved credentials found.", file=sys.stderr)
+        return
+
+    print("Found Credentials for:", file=sys.stderr)
+    for ep in config.credentials.keys():
+        print(" - ", file=sys.stderr, end="", flush=True)
+        print(ep, file=sys.stdout, flush=True)

kleinkram/commands/list.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+from typing import List
+from typing import Optional
+
+import typer
+from kleinkram.api.client import AuthenticatedClient
+from kleinkram.api.routes import get_files
+from kleinkram.api.routes import get_missions
+from kleinkram.api.routes import get_projects
+from kleinkram.config import get_shared_state
+from kleinkram.models import files_to_table
+from kleinkram.models import missions_to_table
+from kleinkram.models import projects_to_table
+from rich.console import Console
+from typer import BadParameter
+
+HELP = """\
+List projects, missions, or files.
+"""
+
+
+list_typer = typer.Typer(
+    name="list", invoke_without_command=True, help=HELP, no_args_is_help=True
+)
+
+
+def _parse_metadata(raw: List[str]) -> dict:
+    ret = {}
+    for tag in raw:
+        if "=" not in tag:
+            raise BadParameter("tag must be formatted as `key=value`")
+        k, v = tag.split("=")
+        ret[k] = v
+    return ret
+
+
+@list_typer.command()
+def files(
+    project: Optional[str] = typer.Option(None, "--project", "-p", help="project name"),
+    mission: Optional[str] = typer.Option(None, "--mission", "-m", help="mission name"),
+    topics: List[str] = typer.Option(None, "--topics", "-t", help="topics"),
+    metadata: Optional[List[str]] = typer.Argument(None, help="tag=value pairs"),
+) -> None:
+    client = AuthenticatedClient()
+
+    _topics = topics if topics else None
+    _metadata = _parse_metadata(metadata or [])
+
+    files = get_files(
+        client, project=project, mission=mission, tags=_metadata, topics=_topics
+    )
+
+    if get_shared_state().verbose:
+        table = files_to_table(files)
+        console = Console()
+        console.print(table)
+    else:
+        for file in files:
+            print(file.id)
+
+
+@list_typer.command()
+def projects() -> None:
+    client = AuthenticatedClient()
+    projects = get_projects(client)
+
+    if get_shared_state().verbose:
+        table = projects_to_table(projects)
+        console = Console()
+        console.print(table)
+    else:
+        for project in projects:
+            print(project.id)
+
+
+@list_typer.command()
+def missions(
+    project: Optional[str] = typer.Option(None, "--project", "-p", help="project name"),
+    metadata: Optional[List[str]] = typer.Argument(None, help="tag=value pairs"),
+) -> None:
+    client = AuthenticatedClient()
+
+    _metadata = _parse_metadata(metadata or [])
+    missions = get_missions(client, project=project, tags=_metadata)
+
+    if get_shared_state().verbose:
+        table = missions_to_table(missions)
+        console = Console()
+        console.print(table)
+    else:
+        for mission in missions:
+            print(mission.id)

kleinkram/commands/mission.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import typer
+from kleinkram.api.client import AuthenticatedClient
+from kleinkram.api.routes import get_mission_by_spec
+from kleinkram.api.routes import update_mission_metadata
+from kleinkram.errors import MissionDoesNotExist
+from kleinkram.utils import get_valid_mission_spec
+from kleinkram.utils import load_metadata
+from kleinkram.utils import to_name_or_uuid
+
+mission_typer = typer.Typer(
+    no_args_is_help=True, context_settings={"help_option_names": ["-h", "--help"]}
+)
+
+
+UPDATE_HELP = """\
+Update a mission.
+"""
+
+NOT_IMPLEMENTED_YET = "Not implemented yet"
+
+
+@mission_typer.command(help=UPDATE_HELP)
+def update(
+    project: Optional[str] = typer.Option(
+        None, "--project", "-p", help="project id or name"
+    ),
+    mission: str = typer.Option(..., "--mission", "-m", help="mission id or name"),
+    metadata: str = typer.Option(help="path to metadata file (json or yaml)"),
+) -> None:
+    _project = to_name_or_uuid(project) if project else None
+    _mission = to_name_or_uuid(mission) if mission else None
+
+    client = AuthenticatedClient()
+
+    mission_spec = get_valid_mission_spec(_mission, _project)
+    mission_parsed = get_mission_by_spec(client, mission_spec)
+
+    if mission_parsed is None:
+        raise MissionDoesNotExist(f"Mission {mission} does not exist")
+
+    metadata_dct = load_metadata(Path(metadata))
+    update_mission_metadata(client, mission_parsed.id, metadata_dct)
+
+
+@mission_typer.command(help=NOT_IMPLEMENTED_YET)
+def create() -> None:
+    raise NotImplementedError("Not implemented yet")
+
+
+@mission_typer.command(help=NOT_IMPLEMENTED_YET)
+def delete() -> None:
+    raise NotImplementedError("Not implemented yet")

kleinkram/commands/project.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+import typer
+
+project_typer = typer.Typer(
+    no_args_is_help=True, context_settings={"help_option_names": ["-h", "--help"]}
+)
+
+NOT_IMPLEMENTED_YET = "Not implemented yet"
+
+
+@project_typer.command(help=NOT_IMPLEMENTED_YET)
+def update() -> None:
+    raise NotImplementedError("Not implemented yet")
+
+
+@project_typer.command(help=NOT_IMPLEMENTED_YET)
+def create() -> None:
+    raise NotImplementedError("Not implemented yet")
+
+
+@project_typer.command(help=NOT_IMPLEMENTED_YET)
+def delete() -> None:
+    raise NotImplementedError("Not implemented yet")

kleinkram/commands/upload.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+from pathlib import Path
+from typing import List
+from typing import Optional
+from uuid import UUID
+
+import typer
+from kleinkram.api.client import AuthenticatedClient
+from kleinkram.api.file_transfer import upload_files
+from kleinkram.api.routes import create_mission
+from kleinkram.api.routes import get_mission_by_id
+from kleinkram.api.routes import get_mission_by_spec
+from kleinkram.api.routes import get_project_id_by_name
+from kleinkram.api.routes import get_tags_map
+from kleinkram.config import get_shared_state
+from kleinkram.errors import MissionDoesNotExist
+from kleinkram.models import MissionByName
+from kleinkram.utils import check_file_paths
+from kleinkram.utils import get_filename_map
+from kleinkram.utils import get_valid_mission_spec
+from kleinkram.utils import load_metadata
+from kleinkram.utils import to_name_or_uuid
+from rich.console import Console
+
+
+HELP = """\
+Upload files to kleinkram.
+"""
+
+upload_typer = typer.Typer(
+    name="upload",
+    no_args_is_help=True,
+    invoke_without_command=True,
+    help=HELP,
+)
+
+
+@upload_typer.callback()
+def upload(
+    files: List[str] = typer.Argument(help="files to upload"),
+    project: Optional[str] = typer.Option(
+        None, "--project", "-p", help="project id or name"
+    ),
+    mission: str = typer.Option(..., "--mission", "-m", help="mission id or name"),
+    create: bool = typer.Option(False, help="create mission if it does not exist"),
+    metadata: Optional[str] = typer.Option(
+        None, help="path to metadata file (json or yaml)"
+    ),
+    fix_filenames: bool = typer.Option(False, help="fix filenames"),
+    ignore_missing_tags: bool = typer.Option(False, help="ignore mission tags"),
+) -> None:
+    _project = to_name_or_uuid(project) if project else None
+    _mission = to_name_or_uuid(mission)
+
+    client = AuthenticatedClient()
+
+    # check files and `fix` filenames
+    if files is None:
+        files = []
+
+    file_paths = [Path(file) for file in files]
+    check_file_paths(file_paths)
+
+    files_map = get_filename_map(
+        [Path(file) for file in files],
+    )
+
+    if not fix_filenames:
+        for name, path in files_map.items():
+            if name != path.name:
+                raise ValueError(
+                    f"invalid filename format {path.name}, use `--fix-filenames`"
+                )
+
+    # parse the mission spec and get mission
+    mission_spec = get_valid_mission_spec(_mission, _project)
+    mission_parsed = get_mission_by_spec(client, mission_spec)
+
+    if not create and mission_parsed is None:
+        raise MissionDoesNotExist(f"mission: {mission} does not exist, use `--create`")
+
+    # create missing mission
+    if mission_parsed is None:
+        if not isinstance(mission_spec, MissionByName):
+            raise ValueError(
+                "cannot create mission using mission id, pecify a mission name"
+            )
+
+        # get the metadata
+        tags_dct = {}
+        if metadata is not None:
+            metadata_dct = load_metadata(Path(metadata))
+            tags_dct = get_tags_map(client, metadata_dct)
+
+        # get project id
+        if isinstance(mission_spec.project, UUID):
+            project_id = mission_spec.project
+        else:
+            project_id = get_project_id_by_name(client, mission_spec.project)
+            if project_id is None:
+                raise ValueError(
+                    f"unable to create mission, project: {mission_spec.project} not found"
+                )
+
+        mission_id = create_mission(
+            client,
+            project_id,
+            mission_spec.name,
+            tags=tags_dct,
+            ignore_missing_tags=ignore_missing_tags,
+        )
+
+        mission_parsed = get_mission_by_id(client, mission_id)
+        assert mission_parsed is not None, "unreachable"
+
+    filtered_files_map = {}
+    remote_file_names = [file.name for file in mission_parsed.files]
+
+    console = Console()
+    for name, path in files_map.items():
+        if name in remote_file_names:
+            console.print(
+                f"file: {name} (path: {path}) already exists in mission", style="yellow"
+            )
+        else:
+            filtered_files_map[name] = path
+
+    if not filtered_files_map:
+        console.print("\nNO FILES UPLOADED", style="yellow")
+        return
+
+    upload_files(
+        filtered_files_map,
+        mission_parsed.id,
+        n_workers=2,
+        verbose=get_shared_state().verbose,
+    )

kleinkram/commands/verify.py

@@ -0,0 +1,117 @@
+from __future__ import annotations
+
+import sys
+from enum import Enum
+from pathlib import Path
+from typing import List
+from typing import Optional
+
+import typer
+from kleinkram.api.client import AuthenticatedClient
+from kleinkram.api.routes import get_mission_by_spec
+from kleinkram.config import get_shared_state
+from kleinkram.errors import MissionDoesNotExist
+from kleinkram.models import FileState
+from kleinkram.utils import b64_md5
+from kleinkram.utils import check_file_paths
+from kleinkram.utils import get_filename_map
+from kleinkram.utils import get_valid_mission_spec
+from kleinkram.utils import to_name_or_uuid
+from rich.console import Console
+from rich.table import Table
+from rich.text import Text
+from tqdm import tqdm
+
+
+class FileVerificationStatus(str, Enum):
+    UPLAODED = "uploaded"
+    UPLOADING = "uploading"
+    MISSING = "missing"
+    CORRUPTED = "hash mismatch"
+    UNKNOWN = "unknown"
+
+
+FILE_STATUS_STYLES = {
+    FileVerificationStatus.UPLAODED: "green",
+    FileVerificationStatus.UPLOADING: "yellow",
+    FileVerificationStatus.MISSING: "yellow",
+    FileVerificationStatus.CORRUPTED: "red",
+    FileVerificationStatus.UNKNOWN: "gray",
+}
+
+
+HELP = """\
+Verify if files were uploaded correctly.
+"""
+
+verify_typer = typer.Typer(name="verify", invoke_without_command=True, help=HELP)
+
+
+@verify_typer.callback()
+def verify(
+    files: List[str] = typer.Argument(help="files to upload"),
+    project: Optional[str] = typer.Option(
+        None, "--project", "-p", help="project id or name"
+    ),
+    mission: str = typer.Option(..., "--mission", "-m", help="mission id or name"),
+    skip_hash: bool = typer.Option(False, help="skip hash check"),
+) -> None:
+    _project = to_name_or_uuid(project) if project else None
+    _mission = to_name_or_uuid(mission)
+
+    client = AuthenticatedClient()
+
+    if files is None:
+        files = []
+
+    mission_spec = get_valid_mission_spec(_mission, _project)
+    mission_parsed = get_mission_by_spec(client, mission_spec)
+
+    if mission_parsed is None:
+        raise MissionDoesNotExist(f"Mission {mission} does not exist")
+
+    # check file types
+    file_paths = [Path(file) for file in files]
+    check_file_paths(file_paths)
+
+    filename_map = get_filename_map(file_paths)
+    remote_files = {file.name: file for file in mission_parsed.files}
+
+    status_dct = {}
+    for name, file in tqdm(
+        filename_map.items(),
+        desc="verifying files",
+        unit="file",
+        disable=skip_hash or not get_shared_state().verbose,
+    ):
+        if name not in remote_files:
+            status_dct[file] = FileVerificationStatus.MISSING
+            continue
+
+        state = remote_files[name].state
+
+        if state == FileState.UPLOADING:
+            status_dct[file] = FileVerificationStatus.UPLOADING
+        elif state == FileState.OK and remote_files[name].hash != b64_md5(file):
+            status_dct[file] = FileVerificationStatus.CORRUPTED
+        elif state == FileState.OK:
+            status_dct[file] = FileVerificationStatus.UPLAODED
+        else:
+            status_dct[file] = FileVerificationStatus.UNKNOWN
+
+    if get_shared_state().verbose:
+        table = Table(title="file status")
+        table.add_column("filename", style="cyan")
+        table.add_column("status", style="green")
+
+        for path, status in status_dct.items():
+            table.add_row(str(path), Text(status, style=FILE_STATUS_STYLES[status]))
+
+        console = Console()
+        console.print(table)
+    else:
+        for path, status in status_dct.items():
+            stream = (
+                sys.stdout if status == FileVerificationStatus.UPLAODED else sys.stderr
+            )
+            print(path, file=stream)