kscale 0.1.3__tar.gz → 0.1.4__tar.gz

{kscale-0.1.3/kscale.egg-info → kscale-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: kscale
-Version: 0.1.3
+Version: 0.1.4
 Summary: The kscale project
 Home-page: https://github.com/kscalelabs/kscale
 Author: Benjamin Bolte

{kscale-0.1.3 → kscale-0.1.4}/kscale/__init__.py

@@ -1,6 +1,6 @@
 """Defines the common interface for the K-Scale Python API."""
 
-__version__ = "0.1.3"
+__version__ = "0.1.4"
 
 from pathlib import Path
 

{kscale-0.1.3 → kscale-0.1.4}/kscale/web/api.py

@@ -12,3 +12,7 @@ class WebAPI(APIBase):
     async def get_profile_info(self) -> ProfileResponse:
         client = await self.www_client()
         return await client.get_profile_info()
+
+    async def get_api_key(self, num_hours: int = 24) -> str:
+        client = await self.www_client()
+        return await client.get_api_key(num_hours)

kscale-0.1.4/kscale/web/cli/user.py

@@ -0,0 +1,53 @@
+"""Defines the CLI for getting information about the current user."""
+
+import logging
+
+import click
+from tabulate import tabulate
+
+from kscale.utils.cli import coro
+from kscale.web.clients.user import UserClient
+
+logger = logging.getLogger(__name__)
+
+
+@click.group()
+def cli() -> None:
+    """Get information about the currently-authenticated user."""
+    pass
+
+
+@cli.command()
+@coro
+async def me() -> None:
+    """Get information about the currently-authenticated user."""
+    client = UserClient()
+    profile = await client.get_profile_info()
+    click.echo(
+        tabulate(
+            [
+                ["Email", profile.email],
+                ["Email verified", profile.email_verified],
+                ["User ID", profile.user.user_id],
+                ["Is admin", profile.user.is_admin],
+                ["Can upload", profile.user.can_upload],
+                ["Can test", profile.user.can_test],
+            ],
+            headers=["Key", "Value"],
+            tablefmt="simple",
+        )
+    )
+
+
+@cli.command()
+@coro
+async def key() -> None:
+    """Get an API key for the currently-authenticated user."""
+    client = UserClient()
+    api_key = await client.get_api_key()
+    click.echo("API key:")
+    click.echo(click.style(api_key, fg="green"))
+
+
+if __name__ == "__main__":
+    cli()

{kscale-0.1.3 → kscale-0.1.4}/kscale/web/clients/base.py

@@ -3,6 +3,7 @@
 import asyncio
 import json
 import logging
+import os
 import secrets
 import time
 import webbrowser
@@ -19,13 +20,16 @@ from pydantic import BaseModel
 from yarl import URL
 
 from kscale.web.gen.api import OICDInfo
-from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_root, get_cache_dir
+from kscale.web.utils import DEFAULT_UPLOAD_TIMEOUT, get_api_root, get_auth_dir
 
 logger = logging.getLogger(__name__)
 
 # This port matches the available port for the OAuth callback.
 OAUTH_PORT = 16821
 
+# This is the name of the API key header for the K-Scale WWW API.
+HEADER_NAME = "x-kscale-api-key"
+
 
 class OAuthCallback:
     def __init__(self) -> None:
@@ -162,7 +166,7 @@ class BaseClient:
 
     @alru_cache
     async def _get_oicd_info(self) -> OICDInfo:
-        cache_path = get_cache_dir() / "oicd_info.json"
+        cache_path = get_auth_dir() / "oicd_info.json"
         if self.use_cache and cache_path.exists():
             with open(cache_path, "r") as f:
                 return OICDInfo(**json.load(f))
@@ -180,7 +184,7 @@ class BaseClient:
         Returns:
             The OpenID Connect server configuration.
         """
-        cache_path = get_cache_dir() / "oicd_metadata.json"
+        cache_path = get_auth_dir() / "oicd_metadata.json"
         if self.use_cache and cache_path.exists():
             with open(cache_path, "r") as f:
                 return json.load(f)
@@ -207,7 +211,7 @@ class BaseClient:
         auth_endpoint = metadata["authorization_endpoint"]
 
         # Use the cached state and nonce if available, otherwise generate.
-        state_file = get_cache_dir() / "oauth_state.json"
+        state_file = get_auth_dir() / "oauth_state.json"
         state: str | None = None
         nonce: str | None = None
         if state_file.exists():
@@ -301,7 +305,7 @@ class BaseClient:
         Returns:
             A bearer token to use with the K-Scale WWW API.
         """
-        cache_path = get_cache_dir() / "bearer_token.txt"
+        cache_path = get_auth_dir() / "bearer_token.txt"
         if self.use_cache and cache_path.exists():
             token = cache_path.read_text()
             if not await self._is_token_expired(token):
@@ -315,9 +319,16 @@ class BaseClient:
     async def get_client(self, *, auth: bool = True) -> httpx.AsyncClient:
         client = self._client if auth else self._client_no_auth
         if client is None:
+            headers: dict[str, str] = {}
+            if auth:
+                if "KSCALE_API_KEY" in os.environ:
+                    headers[HEADER_NAME] = os.environ["KSCALE_API_KEY"]
+                else:
+                    headers["Authorization"] = f"Bearer {await self.get_bearer_token()}"
+
             client = httpx.AsyncClient(
                 base_url=self.base_url,
-                headers={"Authorization": f"Bearer {await self.get_bearer_token()}"} if auth else None,
+                headers=headers,
                 timeout=httpx.Timeout(30.0),
             )
             if auth:

{kscale-0.1.3 → kscale-0.1.4}/kscale/web/clients/robot_class.py

@@ -3,6 +3,7 @@
 import hashlib
 import json
 import logging
+import tarfile
 from pathlib import Path
 
 import httpx
@@ -13,7 +14,7 @@ from kscale.web.gen.api import (
     RobotDownloadURDFResponse,
     RobotUploadURDFResponse,
 )
-from kscale.web.utils import get_cache_dir, should_refresh_file
+from kscale.web.utils import get_robots_dir, should_refresh_file
 
 logger = logging.getLogger(__name__)
 
@@ -98,7 +99,7 @@ class RobotClassClient(BaseClient):
         return response
 
     async def download_robot_class_urdf(self, class_name: str, *, cache: bool = True) -> Path:
-        cache_path = get_cache_dir() / class_name / "robot.tgz"
+        cache_path = get_robots_dir() / class_name / "robot.tgz"
         if cache and cache_path.exists() and not should_refresh_file(cache_path):
             return cache_path
         data = await self._request("GET", f"/robots/urdf/{class_name}", auth=True)
@@ -132,7 +133,10 @@ class RobotClassClient(BaseClient):
         if hash_value_hex != expected_hash:
             raise ValueError(f"MD5 hash mismatch: {hash_value_hex} != {expected_hash}")
 
-        # Updates the info file.
+        logger.info("Unpacking downloaded file")
+        with tarfile.open(cache_path, "r:gz") as tar:
+            tar.extractall(path=cache_path.parent)
+
         logger.info("Updating downloaded file information")
         info = {"md5_hash": hash_value_hex}
         with open(cache_path_info, "w") as f:

{kscale-0.1.3 → kscale-0.1.4}/kscale/web/clients/user.py

@@ -8,3 +8,7 @@ class UserClient(BaseClient):
     async def get_profile_info(self) -> ProfileResponse:
         data = await self._request("GET", "/auth/profile", auth=True)
         return ProfileResponse(**data)
+
+    async def get_api_key(self, num_hours: int = 24) -> str:
+        data = await self._request("POST", "/auth/key", auth=True, data={"num_hours": num_hours})
+        return data["api_key"]

{kscale-0.1.3 → kscale-0.1.4}/kscale/web/utils.py

@@ -13,24 +13,26 @@ DEFAULT_UPLOAD_TIMEOUT = 300.0  # 5 minutes
 
 
 @functools.lru_cache
-def get_cache_dir() -> Path:
+def get_kscale_dir() -> Path:
     """Returns the cache directory for artifacts."""
     return Path(Settings.load().www.cache_dir).expanduser().resolve()
 
 
+def get_auth_dir() -> Path:
+    """Returns the directory for authentication artifacts."""
+    return get_kscale_dir() / "auth"
+
+
+def get_robots_dir() -> Path:
+    """Returns the directory for robot artifacts."""
+    return get_kscale_dir() / "robots"
+
+
 def should_refresh_file(file: Path) -> bool:
     """Returns whether the file should be refreshed."""
     return file.exists() and file.stat().st_mtime < time.time() - Settings.load().www.refresh_interval_minutes * 60
 
 
-@functools.lru_cache
-def get_artifact_dir(artifact_id: str) -> Path:
-    """Returns the directory for a specific artifact."""
-    cache_dir = get_cache_dir() / artifact_id
-    cache_dir.mkdir(parents=True, exist_ok=True)
-    return cache_dir
-
-
 @functools.lru_cache
 def get_api_root() -> str:
     """Returns the root URL for the K-Scale WWW API."""

{kscale-0.1.3 → kscale-0.1.4/kscale.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: kscale
-Version: 0.1.3
+Version: 0.1.4
 Summary: The kscale project
 Home-page: https://github.com/kscalelabs/kscale
 Author: Benjamin Bolte

kscale-0.1.3/kscale/web/cli/user.py

@@ -1,33 +0,0 @@
-"""Defines the CLI for getting information about the current user."""
-
-import logging
-
-import click
-
-from kscale.utils.cli import coro
-from kscale.web.clients.user import UserClient
-
-logger = logging.getLogger(__name__)
-
-
-@click.group()
-def cli() -> None:
-    """Get information about the currently-authenticated user."""
-    pass
-
-
-@cli.command()
-@coro
-async def me() -> None:
-    client = UserClient()
-    profile = await client.get_profile_info()
-    logger.info("Email: %s", profile.email)
-    logger.info("Email verified: %s", profile.email_verified)
-    logger.info("User ID: %s", profile.user.user_id)
-    logger.info("Is admin: %s", profile.user.is_admin)
-    logger.info("Can upload: %s", profile.user.can_upload)
-    logger.info("Can test: %s", profile.user.can_test)
-
-
-if __name__ == "__main__":
-    cli()