unitlab 2.4.3__tar.gz → 2.4.4__tar.gz

{unitlab-2.4.3 → unitlab-2.4.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: unitlab
-Version: 2.4.3
+Version: 2.4.4
 Summary: Python SDK for the Unitlab.ai data annotation platform
 Project-URL: Homepage, https://unitlab.ai
 Project-URL: Documentation, https://docs.unitlab.ai

{unitlab-2.4.3 → unitlab-2.4.4}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "unitlab"
-version = "2.4.3"
+version = "2.4.4"
 description = "Python SDK for the Unitlab.ai data annotation platform"
 readme = "README.md"
 license = "MIT"

{unitlab-2.4.3 → unitlab-2.4.4}/src/unitlab/cli.py

@@ -37,12 +37,6 @@ class DownloadType(str, Enum):
     files = "files"
 
 
-class DatasetSplit(str, Enum):
-    train = "train"
-    validation = "validation"
-    test = "test"
-
-
 @app.command(help="Configure the credentials")
 def configure(
     api_key: Annotated[
@@ -124,11 +118,14 @@ def dataset_download(
         typer.Option(help="Export type (e.g. COCO, YOLOv8, YOLOv5, UUEF)"),
     ] = "COCO",
     split_type: Annotated[
-        DatasetSplit, typer.Option(help="Dataset split (train, validation, test)")
-    ] = DatasetSplit.train,
+        str | None,
+        typer.Option(
+            help="Dataset split (e.g. train, validation). Auto-detected if omitted."
+        ),
+    ] = None,
 ):
     if download_type == DownloadType.annotation:
-        get_client(api_key).dataset_download(str(pk), export_type, split_type.value)
+        get_client(api_key).dataset_download(str(pk), export_type, split_type)
     else:
         typer.echo(get_client(api_key).dataset_download_files(str(pk)))
 

{unitlab-2.4.3 → unitlab-2.4.4}/src/unitlab/client.py

@@ -308,9 +308,41 @@ class UnitlabClient:
     def datasets(self, pretty: int = 0) -> Any:
         return self._get(f"/api/sdk/datasets/?pretty={pretty}")
 
+    def dataset_info(self, dataset_id: str) -> Any:
+        """Return metadata for a dataset.
+
+        Includes name, version, annotation_type, number_of_data,
+        download_formats and the list of available ``splits``.
+        """
+        return self._get(f"/api/sdk/datasets/{dataset_id}/")
+
     def dataset_download(
-        self, dataset_id: str, export_type: str, split_type: str
+        self,
+        dataset_id: str,
+        export_type: str,
+        split_type: str | None = None,
     ) -> str:
+        if split_type is None:
+            info = self.dataset_info(dataset_id)
+            splits = info.get("splits") or []
+            if not splits:
+                raise NetworkError(
+                    message=(
+                        "No splits available for this dataset. Nothing to download."
+                    ),
+                    detail=None,
+                )
+            if len(splits) > 1:
+                raise NetworkError(
+                    message=(
+                        f"Multiple splits available: {', '.join(splits)}. "
+                        f"Specify one with split_type=..."
+                    ),
+                    detail=None,
+                )
+            split_type = splits[0]
+            logger.info(f"Auto-selected split: {split_type}")
+
         response = self._post(
             f"/api/sdk/datasets/{dataset_id}/",
             data={