diff --git a/kscale/store/client.py b/kscale/store/client.py index aef18a7..7e7a52f 100644 --- a/kscale/store/client.py +++ b/kscale/store/client.py @@ -26,6 +26,7 @@ def __init__(self, base_url: str = get_api_root()) -> None: self.client = httpx.AsyncClient( base_url=self.base_url, headers={"Authorization": f"Bearer {get_api_key()}"}, + timeout=httpx.Timeout(30.0), ) async def _request( diff --git a/kscale/store/urdf.py b/kscale/store/urdf.py index 940d87b..6824610 100644 --- a/kscale/store/urdf.py +++ b/kscale/store/urdf.py @@ -21,6 +21,17 @@ 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() @@ -75,7 +86,7 @@ def create_tarball(folder_path: Path, output_filename: str, cache_dir: 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 (".urdf", ".mjcf", ".stl", ".obj", ".dae"): + 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: