kleinkram 0.38.1.dev20241212075157__py3-none-any.whl → 0.38.1.dev20250113080249__py3-none-any.whl

kleinkram/cli/error_handling.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from collections import OrderedDict
+from typing import Any
+from typing import Callable
+from typing import Type
+
+import typer
+from click import ClickException
+
+ExceptionHandler = Callable[[Exception], int]
+
+
+class ErrorHandledTyper(typer.Typer):
+    """\
+    error handlers that are last added will be used first
+    """
+
+    _error_handlers: OrderedDict[Type[Exception], ExceptionHandler]
+
+    def error_handler(
+        self, exc: Type[Exception]
+    ) -> Callable[[ExceptionHandler], ExceptionHandler]:
+        def dec(func: ExceptionHandler) -> ExceptionHandler:
+            self._error_handlers[exc] = func
+            return func
+
+        return dec
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._error_handlers = OrderedDict()
+
+    def __call__(self, *args: Any, **kwargs: Any) -> int:
+        try:
+            return super().__call__(*args, **kwargs)
+        except Exception as e:
+            if isinstance(e, ClickException):
+                raise
+            for tp, handler in reversed(self._error_handlers.items()):
+                if isinstance(e, tp):
+                    exit_code = handler(e)
+                    raise SystemExit(exit_code)
+            raise

kleinkram/config.py

@@ -1,18 +1,31 @@
+"""
+this file contains a global config and a global state object
+
+to get the config use `get_config()`
+"""
+
 from __future__ import annotations
 
 import json
+import logging
 import os
 import tempfile
 from dataclasses import dataclass
+from dataclasses import field
 from enum import Enum
 from pathlib import Path
+from typing import Any
 from typing import Dict
 from typing import NamedTuple
 from typing import Optional
 
+from rich.table import Table
+from rich.text import Text
+
 from kleinkram._version import __local__
 from kleinkram._version import __version__
-from kleinkram.errors import InvalidConfigFile
+
+logger = logging.getLogger(__name__)
 
 CONFIG_PATH = Path().home() / ".kleinkram.json"
 
@@ -23,13 +36,33 @@ class Environment(Enum):
     PROD = "prod"
 
 
-DEFAULT_API = {
-    Environment.LOCAL: "http://localhost:3000",
-    Environment.DEV: "https://api.datasets.dev.leggedrobotics.com",
-    Environment.PROD: "https://api.datasets.leggedrobotics.com",
-}
+class Endpoint(NamedTuple):
+    name: str
+    api: str
+    s3: str
+
+
+class Credentials(NamedTuple):
+    auth_token: Optional[str] = None
+    refresh_token: Optional[str] = None
+    cli_key: Optional[str] = None
+
+
+DEFAULT_LOCAL_API = "http://localhost:3000"
+DEFAULT_LOCAL_S3 = "http://localhost:9000"
+
+DEFAULT_DEV_API = "https://api.datasets.dev.leggedrobotics.com"
+DEFAULT_DEV_S3 = "https://s3.datasets.dev.leggedrobotics.com"
+
+DEFAULT_PROD_API = "https://api.datasets.leggedrobotics.com"
+DEFAULT_PROD_S3 = "https://s3.datasets.leggedrobotics.com"
 
-LOCAL_S3 = "http://localhost:9000"
+
+DEFAULT_ENDPOINTS = {
+    "local": Endpoint("local", DEFAULT_LOCAL_API, DEFAULT_LOCAL_S3),
+    "dev": Endpoint("dev", DEFAULT_DEV_API, DEFAULT_DEV_S3),
+    "prod": Endpoint("prod", DEFAULT_PROD_API, DEFAULT_PROD_S3),
+}
 
 
 def get_env() -> Environment:
@@ -40,131 +73,132 @@ def get_env() -> Environment:
     return Environment.PROD
 
 
-def get_default_endpoints() -> str:
-    env = get_env()
-    return DEFAULT_API[env]
+@dataclass
+class Config:
+    version: str = __version__
+    selected_endpoint: str = field(default_factory=lambda: get_env().value)
+    endpoints: Dict[str, Endpoint] = field(
+        default_factory=lambda: DEFAULT_ENDPOINTS.copy()
+    )
+    endpoint_credentials: Dict[str, Credentials] = field(default_factory=dict)
+
+    @property
+    def endpoint(self) -> Endpoint:
+        return self.endpoints[self.selected_endpoint]
+
+    @endpoint.setter
+    def endpoint(self, value: Endpoint) -> None:
+        self.endpoints[self.selected_endpoint] = value
 
+    @property
+    def credentials(self) -> Optional[Credentials]:
+        return self.endpoint_credentials.get(self.selected_endpoint)
 
-class Credentials(NamedTuple):
-    auth_token: Optional[str] = None
-    refresh_token: Optional[str] = None
-    cli_key: Optional[str] = None
+    @credentials.setter
+    def credentials(self, value: Credentials) -> None:
+        self.endpoint_credentials[self.selected_endpoint] = value
 
 
-JSON_ENDPOINT_KEY = "endpoint"
-JSON_CREDENTIALS_KEY = "credentials"
+def _config_to_dict(config: Config) -> Dict[str, Any]:
+    return {
+        "version": config.version,
+        "endpoints": {key: value._asdict() for key, value in config.endpoints.items()},
+        "endpoint_credentials": {
+            key: value._asdict() for key, value in config.endpoint_credentials.items()
+        },
+        "selected_endpoint": config.endpoint.name,
+    }
 
 
-class Config:
-    endpoint: str
-    credentials: Dict[str, Credentials]
-
-    def __init__(self, overwrite: bool = False) -> None:
-        default_endpoint = get_default_endpoints()
-
-        self.credentials = {}
-        self.endpoint = default_endpoint
-
-        if not CONFIG_PATH.exists():
-            self.save()
-
-        try:
-            self._read_config()
-        except InvalidConfigFile:
-            if not overwrite:
-                self.credentials = {}
-                self.endpoint = default_endpoint
-                self.save()
-            else:
-                raise
-
-    def _read_config(self) -> None:
-        with open(CONFIG_PATH, "r") as file:
-            try:
-                content = json.load(file)
-            except Exception:
-                raise InvalidConfigFile
-
-        endpoint = content.get(JSON_ENDPOINT_KEY, None)
-        if not isinstance(endpoint, str):
-            raise InvalidConfigFile
-
-        credentials = content.get(JSON_CREDENTIALS_KEY, None)
-        if not isinstance(credentials, dict):
-            raise InvalidConfigFile
-
-        try:
-            parsed_creds = {}
-            for ep, creds in credentials.items():
-                parsed_creds[ep] = Credentials(**creds)
-        except Exception:
-            raise InvalidConfigFile
-
-        self.endpoint = endpoint
-        self.credentials = parsed_creds
+def _config_from_dict(dct: Dict[str, Any]) -> Config:
+    return Config(
+        dct["version"],
+        dct["selected_endpoint"],
+        {key: Endpoint(**value) for key, value in dct["endpoints"].items()},
+        {
+            key: Credentials(**value)
+            for key, value in dct["endpoint_credentials"].items()
+        },
+    )
 
-    @property
-    def has_cli_key(self) -> bool:
-        if self.endpoint not in self.credentials:
-            return False
-        return self.credentials[self.endpoint].cli_key is not None
 
-    @property
-    def has_refresh_token(self) -> bool:
-        if self.endpoint not in self.credentials:
-            return False
-        return self.credentials[self.endpoint].refresh_token is not None
+def save_config(config: Config, path: Path = CONFIG_PATH) -> None:
+    fd, temp_path = tempfile.mkstemp()
+    with os.fdopen(fd, "w") as f:
+        json.dump(_config_to_dict(config), f)
+    os.replace(temp_path, path)
 
-    @property
-    def auth_token(self) -> Optional[str]:
-        return self.credentials[self.endpoint].auth_token
 
-    @property
-    def refresh_token(self) -> Optional[str]:
-        return self.credentials[self.endpoint].refresh_token
+def _load_config(*, path: Path = CONFIG_PATH) -> Config:
+    if not path.exists():
+        return Config()
+    with open(path, "r") as f:
+        return _config_from_dict(json.load(f))
 
-    @property
-    def cli_key(self) -> Optional[str]:
-        return self.credentials[self.endpoint].cli_key
 
-    def save(self) -> None:
-        serialized_tokens = {}
-        for endpoint, auth in self.credentials.items():
-            serialized_tokens[endpoint] = auth._asdict()
+LOADED_CONFIGS: Dict[Path, Config] = {}
+
+
+def get_config(path: Path = CONFIG_PATH) -> Config:
+    if path not in LOADED_CONFIGS:
+        LOADED_CONFIGS[path] = _load_config(path=path)
+    return LOADED_CONFIGS[path]
+
+
+def select_endpoint(config: Config, name: str, path: Path = CONFIG_PATH) -> None:
+    if name not in config.endpoints:
+        raise ValueError(f"Endpoint {name} not found.")
+    config.selected_endpoint = name
+    save_config(config, path)
+
+
+def add_endpoint(config: Config, endpoint: Endpoint, path: Path = CONFIG_PATH) -> None:
+    config.endpoints[endpoint.name] = endpoint
+    config.selected_endpoint = endpoint.name
+    save_config(config, path)
+
 
-        data = {
-            JSON_ENDPOINT_KEY: self.endpoint,
-            JSON_CREDENTIALS_KEY: serialized_tokens,
-        }
+def check_config_compatibility(path: Path = CONFIG_PATH) -> bool:
+    """\
+    returns `False` if config file exists but is not compatible with the current version
 
-        # atomically write to file
-        fd, tmp_path = tempfile.mkstemp()
-        with open(fd, "w") as file:
-            json.dump(data, file)
+    TODO: add more sophisticated version checking
+    """
+    if not path.exists():
+        return True
+    try:
+        _ = _load_config(path=path)
+    except Exception as e:
+        logger.info(f"Error loading config: {e}")
+        return False
+    return True
 
-        os.replace(tmp_path, CONFIG_PATH)
 
-    def clear_credentials(self, all: bool = False) -> None:
-        if all:
-            self.credentials = {}
-        elif self.endpoint in self.credentials:
-            del self.credentials[self.endpoint]
-        self.save()
+def endpoint_table(config: Config) -> Table:
+    table = Table(title="Available Endpoints")
+    table.add_column("Name", style="cyan")
+    table.add_column("API", style="cyan")
+    table.add_column("S3", style="cyan")
 
-    def save_credentials(self, creds: Credentials) -> None:
-        self.credentials[self.endpoint] = creds
-        self.save()
+    for name, endpoint in config.endpoints.items():
+        display_name = (
+            Text(name, style="bold yellow")
+            if name == config.selected_endpoint
+            else Text(name)
+        )
+        table.add_row(display_name, endpoint.api, endpoint.s3)
+    return table
 
 
 @dataclass
-class _SharedState:
+class SharedState:
     log_file: Optional[Path] = None
     verbose: bool = True
     debug: bool = False
 
 
-SHARED_STATE = _SharedState()
+SHARED_STATE = SharedState()
 
 
-def get_shared_state() -> _SharedState:
+def get_shared_state() -> SharedState:
     return SHARED_STATE

kleinkram/core.py

@@ -1,14 +1,262 @@
+"""
+this file contains the main functionality of kleinkram cli
+
+- download
+- upload
+- verify
+- update_file
+- update_mission
+- update_project
+- delete_files
+- delete_mission
+- delete_project
+"""
+
 from __future__ import annotations
 
 from pathlib import Path
+from typing import Collection
+from typing import Dict
 from typing import List
+from typing import Optional
+from typing import Sequence
 from uuid import UUID
 
+from rich.console import Console
+from tqdm import tqdm
+
+import kleinkram.api.file_transfer
+import kleinkram.api.routes
+import kleinkram.errors
+from kleinkram.api.client import AuthenticatedClient
+from kleinkram.api.query import FileQuery
+from kleinkram.api.query import MissionQuery
+from kleinkram.api.query import ProjectQuery
+from kleinkram.api.query import check_mission_query_is_creatable
+from kleinkram.errors import MissionNotFound
+from kleinkram.models import FileState
+from kleinkram.models import FileVerificationStatus
+from kleinkram.printing import files_to_table
+from kleinkram.utils import b64_md5
+from kleinkram.utils import check_file_paths
+from kleinkram.utils import file_paths_from_files
+from kleinkram.utils import get_filename_map
+
+
+def download(
+    *,
+    client: AuthenticatedClient,
+    query: FileQuery,
+    base_dir: Path,
+    nested: bool = False,
+    overwrite: bool = False,
+    verbose: bool = False,
+) -> None:
+    """\
+    downloads files, asserts that the destition dir exists
+    returns the files that were downloaded
+
+    TODO: the above is a lie, at the moment we just return all files that were found
+    this might include some files that were skipped or not downloaded for some reason
+    we would need to modify the `download_files` function to return this in the future
+    """
+
+    if not base_dir.exists():
+        raise ValueError(f"Destination {base_dir.absolute()} does not exist")
+    if not base_dir.is_dir():
+        raise ValueError(f"Destination {base_dir.absolute()} is not a directory")
+
+    # retrive files and get the destination paths
+    files = list(kleinkram.api.routes.get_files(client, file_query=query))
+    paths = file_paths_from_files(files, dest=base_dir, allow_nested=nested)
+
+    if verbose:
+        table = files_to_table(files, title="downloading files...")
+        Console().print(table)
+
+    kleinkram.api.file_transfer.download_files(
+        client, paths, verbose=verbose, overwrite=overwrite
+    )
+
+
+def upload(
+    *,
+    client: AuthenticatedClient,
+    query: MissionQuery,
+    file_paths: Sequence[Path],
+    create: bool = False,
+    metadata: Optional[Dict[str, str]] = None,
+    ignore_missing_metadata: bool = False,
+    verbose: bool = False,
+) -> None:
+    """\
+    uploads files to a mission
+
+    create a mission if it does not exist if `create` is True
+    in that case you can also specify `metadata` and `ignore_missing_metadata`
+    """
+    # check that file paths are for valid files and have valid suffixes
+    check_file_paths(file_paths)
+
+    try:
+        mission = kleinkram.api.routes.get_mission(client, query=query)
+    except MissionNotFound:
+        if not create:
+            raise
+        mission = None
+
+    if create and mission is None:
+        # check if project exists and get its id at the same time
+        project_id = kleinkram.api.routes.get_project(
+            client, query=query.project_query
+        ).id
+        mission_name = check_mission_query_is_creatable(query)
+        kleinkram.api.routes._create_mission(
+            client,
+            project_id,
+            mission_name,
+            metadata=metadata or {},
+            ignore_missing_tags=ignore_missing_metadata,
+        )
+        mission = kleinkram.api.routes.get_mission(client, query)
+
+    assert mission is not None, "unreachable"
+
+    filename_map = get_filename_map(file_paths)
+    kleinkram.api.file_transfer.upload_files(
+        client, filename_map, mission.id, verbose=verbose
+    )
+
+
+def verify(
+    *,
+    client: AuthenticatedClient,
+    query: MissionQuery,
+    file_paths: Sequence[Path],
+    skip_hash: bool = False,
+    verbose: bool = False,
+) -> Dict[Path, FileVerificationStatus]:
+    # check that file paths are for valid files and have valid suffixes
+    check_file_paths(file_paths)
+
+    # check that the mission exists
+    _ = kleinkram.api.routes.get_mission(client, query)
+
+    remote_files = {
+        f.name: f
+        for f in kleinkram.api.routes.get_files(
+            client, file_query=FileQuery(mission_query=query)
+        )
+    }
+    filename_map = get_filename_map(file_paths)
+
+    # verify files
+    file_status: Dict[Path, FileVerificationStatus] = {}
+    for name, file in tqdm(
+        filename_map.items(),
+        desc="verifying files",
+        unit="file",
+        disable=not verbose,
+    ):
+        if name not in remote_files:
+            file_status[file] = FileVerificationStatus.MISSING
+            continue
+
+        remote_file = remote_files[name]
+
+        if remote_file.state == FileState.UPLOADING:
+            file_status[file] = FileVerificationStatus.UPLOADING
+        elif remote_file.state == FileState.OK:
+            if remote_file.hash is None:
+                file_status[file] = FileVerificationStatus.COMPUTING_HASH
+            elif skip_hash or remote_file.hash == b64_md5(file):
+                file_status[file] = FileVerificationStatus.UPLAODED
+            else:
+                file_status[file] = FileVerificationStatus.MISMATCHED_HASH
+        else:
+            file_status[file] = FileVerificationStatus.UNKNOWN
+    return file_status
+
+
+def update_file(*, client: AuthenticatedClient, file_id: UUID) -> None:
+    """\
+    TODO: what should this even do
+    """
+    _ = client, file_id
+    raise NotImplementedError("if you have an idea what this should do, open an issue")
+
+
+def update_mission(
+    *, client: AuthenticatedClient, mission_id: UUID, metadata: Dict[str, str]
+) -> None:
+    # TODO: this funciton will do more than just overwirte the metadata in the future
+    kleinkram.api.routes._update_mission(client, mission_id, metadata=metadata)
+
+
+def update_project(
+    *,
+    client: AuthenticatedClient,
+    project_id: UUID,
+    description: Optional[str] = None,
+    new_name: Optional[str] = None,
+) -> None:
+    # TODO: this function should do more in the future
+    kleinkram.api.routes._update_project(
+        client, project_id, description=description, new_name=new_name
+    )
+
+
+def delete_files(*, client: AuthenticatedClient, file_ids: Collection[UUID]) -> None:
+    """\
+    deletes multiple files accross multiple missions
+    """
+    files = list(kleinkram.api.routes.get_files(client, FileQuery(ids=list(file_ids))))
+
+    # check if all file_ids were actually found
+    found_ids = [f.id for f in files]
+    for file_id in file_ids:
+        if file_id not in found_ids:
+            raise kleinkram.errors.FileNotFound(
+                f"file {file_id} not found, did not delete any files"
+            )
+
+    # we can only batch delete files within the same mission
+    missions_to_files: Dict[UUID, List[UUID]] = {}
+    for file in files:
+        if file.mission_id not in missions_to_files:
+            missions_to_files[file.mission_id] = []
+        missions_to_files[file.mission_id].append(file.id)
+
+    for mission_id, ids_ in missions_to_files.items():
+        kleinkram.api.routes._delete_files(client, file_ids=ids_, mission_id=mission_id)
+
+
+def delete_mission(*, client: AuthenticatedClient, mission_id: UUID) -> None:
+    mquery = MissionQuery(ids=[mission_id])
+    mission = kleinkram.api.routes.get_mission(client, mquery)
+    files = list(
+        kleinkram.api.routes.get_files(
+            client, file_query=FileQuery(mission_query=mquery)
+        )
+    )
 
-def upload() -> None: ...
+    # delete the files and then the mission
+    kleinkram.api.routes._delete_files(client, [f.id for f in files], mission.id)
+    kleinkram.api.routes._delete_mission(client, mission_id)
 
 
-def download() -> None: ...
+def delete_project(*, client: AuthenticatedClient, project_id: UUID) -> None:
+    pquery = ProjectQuery(ids=[project_id])
+    _ = kleinkram.api.routes.get_project(client, pquery)  # check if project exists
 
+    # delete all missions and files
+    missions = list(
+        kleinkram.api.routes.get_missions(
+            client, mission_query=MissionQuery(project_query=pquery)
+        )
+    )
+    for mission in missions:
+        delete_mission(client=client, mission_id=mission.id)
 
-def download_file(ids: List[UUID], dest: Path) -> None: ...
+    # delete the project
+    kleinkram.api.routes._delete_project(client, project_id)