unitlab 2.4.5__tar.gz → 2.5.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: unitlab
3
- Version: 2.4.5
3
+ Version: 2.5.0
4
4
  Summary: Python SDK for the Unitlab.ai data annotation platform
5
5
  Project-URL: Homepage, https://unitlab.ai
6
6
  Project-URL: Documentation, https://docs.unitlab.ai
@@ -121,8 +121,6 @@ projects = client.projects()
121
121
  # Get project details
122
122
  project = client.project("PROJECT_ID")
123
123
 
124
- # Get project members
125
- members = client.project_members("PROJECT_ID")
126
124
  ```
127
125
 
128
126
  ### Upload data
@@ -134,14 +132,21 @@ client.project_upload_data(
134
132
  )
135
133
  ```
136
134
 
137
- Additional options for specific project types:
135
+ The upload command accepts mixed directories. File type is detected per file from
136
+ the supported extensions: images, text, video, audio, medical
137
+ (`.dcm`, `.nii`, `.nii.gz`, `.nrrd`), and documents (`.pdf`).
138
138
 
139
- ```python
140
- # Text projects
141
- client.project_upload_data("PROJECT_ID", "./docs", sentences_per_chunk=10)
139
+ Additional options for specific file types:
142
140
 
143
- # Video projects
141
+ ```python
142
+ # Video files
144
143
  client.project_upload_data("PROJECT_ID", "./videos", fps=30.0)
144
+
145
+ # PDF documents
146
+ client.project_upload_data("PROJECT_ID", "./pdfs")
147
+
148
+ # DICOM, NIfTI, and NRRD medical volumes
149
+ client.project_upload_data("PROJECT_ID", "./medical")
145
150
  ```
146
151
 
147
152
  ### Datasets
@@ -150,8 +155,11 @@ client.project_upload_data("PROJECT_ID", "./videos", fps=30.0)
150
155
  # List all datasets
151
156
  datasets = client.datasets()
152
157
 
153
- # Download annotations (COCO, YOLOv8, YOLOv5, etc.)
154
- path = client.dataset_download("DATASET_ID", export_type="COCO", split_type="train")
158
+ # Download one annotation split using the release's fixed export format
159
+ path = client.dataset_download("DATASET_ID", split_type="train")
160
+
161
+ # Omit split_type to download all available splits as one bundle
162
+ path = client.dataset_download("DATASET_ID")
155
163
 
156
164
  # Download raw files
157
165
  folder = client.dataset_download_files("DATASET_ID")
@@ -168,11 +176,12 @@ unitlab project list
168
176
  # Project details
169
177
  unitlab project detail PROJECT_ID
170
178
 
171
- # Project members
172
- unitlab project members PROJECT_ID
173
-
174
179
  # Upload data to a project
175
180
  unitlab project upload PROJECT_ID --directory ./images
181
+
182
+ # Upload PDFs or medical volumes with the same command
183
+ unitlab project upload PROJECT_ID --directory ./pdfs
184
+ unitlab project upload PROJECT_ID --directory ./medical
176
185
  ```
177
186
 
178
187
  ### Datasets
@@ -181,8 +190,11 @@ unitlab project upload PROJECT_ID --directory ./images
181
190
  # List datasets
182
191
  unitlab dataset list
183
192
 
184
- # Download annotations
185
- unitlab dataset download DATASET_ID --export-type COCO --split-type train
193
+ # Download one annotation split using the release's fixed export format
194
+ unitlab dataset download DATASET_ID --split-type train
195
+
196
+ # Download all annotation splits as one bundle
197
+ unitlab dataset download DATASET_ID
186
198
 
187
199
  # Download raw files
188
200
  unitlab dataset download DATASET_ID --download-type files
@@ -87,8 +87,6 @@ projects = client.projects()
87
87
  # Get project details
88
88
  project = client.project("PROJECT_ID")
89
89
 
90
- # Get project members
91
- members = client.project_members("PROJECT_ID")
92
90
  ```
93
91
 
94
92
  ### Upload data
@@ -100,14 +98,21 @@ client.project_upload_data(
100
98
  )
101
99
  ```
102
100
 
103
- Additional options for specific project types:
101
+ The upload command accepts mixed directories. File type is detected per file from
102
+ the supported extensions: images, text, video, audio, medical
103
+ (`.dcm`, `.nii`, `.nii.gz`, `.nrrd`), and documents (`.pdf`).
104
104
 
105
- ```python
106
- # Text projects
107
- client.project_upload_data("PROJECT_ID", "./docs", sentences_per_chunk=10)
105
+ Additional options for specific file types:
108
106
 
109
- # Video projects
107
+ ```python
108
+ # Video files
110
109
  client.project_upload_data("PROJECT_ID", "./videos", fps=30.0)
110
+
111
+ # PDF documents
112
+ client.project_upload_data("PROJECT_ID", "./pdfs")
113
+
114
+ # DICOM, NIfTI, and NRRD medical volumes
115
+ client.project_upload_data("PROJECT_ID", "./medical")
111
116
  ```
112
117
 
113
118
  ### Datasets
@@ -116,8 +121,11 @@ client.project_upload_data("PROJECT_ID", "./videos", fps=30.0)
116
121
  # List all datasets
117
122
  datasets = client.datasets()
118
123
 
119
- # Download annotations (COCO, YOLOv8, YOLOv5, etc.)
120
- path = client.dataset_download("DATASET_ID", export_type="COCO", split_type="train")
124
+ # Download one annotation split using the release's fixed export format
125
+ path = client.dataset_download("DATASET_ID", split_type="train")
126
+
127
+ # Omit split_type to download all available splits as one bundle
128
+ path = client.dataset_download("DATASET_ID")
121
129
 
122
130
  # Download raw files
123
131
  folder = client.dataset_download_files("DATASET_ID")
@@ -134,11 +142,12 @@ unitlab project list
134
142
  # Project details
135
143
  unitlab project detail PROJECT_ID
136
144
 
137
- # Project members
138
- unitlab project members PROJECT_ID
139
-
140
145
  # Upload data to a project
141
146
  unitlab project upload PROJECT_ID --directory ./images
147
+
148
+ # Upload PDFs or medical volumes with the same command
149
+ unitlab project upload PROJECT_ID --directory ./pdfs
150
+ unitlab project upload PROJECT_ID --directory ./medical
142
151
  ```
143
152
 
144
153
  ### Datasets
@@ -147,8 +156,11 @@ unitlab project upload PROJECT_ID --directory ./images
147
156
  # List datasets
148
157
  unitlab dataset list
149
158
 
150
- # Download annotations
151
- unitlab dataset download DATASET_ID --export-type COCO --split-type train
159
+ # Download one annotation split using the release's fixed export format
160
+ unitlab dataset download DATASET_ID --split-type train
161
+
162
+ # Download all annotation splits as one bundle
163
+ unitlab dataset download DATASET_ID
152
164
 
153
165
  # Download raw files
154
166
  unitlab dataset download DATASET_ID --download-type files
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "unitlab"
7
- version = "2.4.5"
7
+ version = "2.5.0"
8
8
  description = "Python SDK for the Unitlab.ai data annotation platform"
9
9
  readme = "README.md"
10
10
  license = "MIT"
@@ -25,9 +25,12 @@ app.add_typer(dataset_app, name="dataset", help="Dataset commands")
25
25
 
26
26
 
27
27
  API_KEY = Annotated[
28
- str,
28
+ str | None,
29
29
  typer.Option(
30
- default_factory=_config.get_api_key, help="The api-key obtained from unitlab.ai"
30
+ help=(
31
+ "The api-key obtained from unitlab.ai. If omitted, reads "
32
+ "UNITLAB_API_KEY, then the configured key."
33
+ )
31
34
  ),
32
35
  ]
33
36
 
@@ -55,36 +58,27 @@ def configure(
55
58
  _config.write_config(api_key=api_key, api_url=api_url)
56
59
 
57
60
 
58
- def get_client(api_key: str) -> UnitlabClient:
61
+ def get_client(api_key: str | None) -> UnitlabClient:
59
62
  return UnitlabClient(api_key=api_key)
60
63
 
61
64
 
62
65
  @project_app.command(name="list", help="Project list")
63
- def project_list(api_key: API_KEY):
66
+ def project_list(api_key: API_KEY = None):
64
67
  typer.echo(get_client(api_key).projects(pretty=1))
65
68
 
66
69
 
67
70
  @project_app.command(name="detail", help="Project detail")
68
- def project_detail(pk: UUID, api_key: API_KEY):
71
+ def project_detail(pk: UUID, api_key: API_KEY = None):
69
72
  typer.echo(get_client(api_key).project(project_id=pk, pretty=1))
70
73
 
71
74
 
72
- @project_app.command(help="Project members")
73
- def members(pk: UUID, api_key: API_KEY):
74
- typer.echo(get_client(api_key).project_members(project_id=pk, pretty=1))
75
-
76
-
77
75
  @project_app.command(help="Upload data")
78
76
  def upload(
79
77
  pk: UUID,
80
- api_key: API_KEY,
81
78
  directory: Annotated[
82
79
  Path, typer.Option(help="Directory containing the data to be uploaded")
83
80
  ],
84
- sentences_per_chunk: Annotated[
85
- int,
86
- typer.Option(help="Sentences per chunk for text projects"),
87
- ] = 10,
81
+ api_key: API_KEY = None,
88
82
  fps: Annotated[
89
83
  float,
90
84
  typer.Option(
@@ -92,40 +86,39 @@ def upload(
92
86
  ),
93
87
  ] = 1.0,
94
88
  ):
95
- get_client(api_key).project_upload_data(
89
+ result = get_client(api_key).project_upload_data(
96
90
  str(pk),
97
91
  directory=directory,
98
- sentences_per_chunk=sentences_per_chunk,
99
92
  fps=fps,
100
93
  )
94
+ typer.echo(f"Uploaded {result['uploaded']} of {result['total']} files")
101
95
 
102
96
 
103
97
  @dataset_app.command(name="list", help="List datasets")
104
- def dataset_list(api_key: API_KEY):
98
+ def dataset_list(api_key: API_KEY = None):
105
99
  typer.echo(get_client(api_key).datasets(pretty=1))
106
100
 
107
101
 
108
102
  @dataset_app.command(name="download", help="Download dataset")
109
103
  def dataset_download(
110
104
  pk: UUID,
111
- api_key: API_KEY,
112
105
  download_type: Annotated[
113
106
  DownloadType,
114
107
  typer.Option(help="Download type (annotation, files)"),
115
108
  ] = DownloadType.annotation,
116
- export_type: Annotated[
117
- str,
118
- typer.Option(help="Export type (e.g. COCO, YOLOv8, YOLOv5, UUEF)"),
119
- ] = "COCO",
109
+ api_key: API_KEY = None,
120
110
  split_type: Annotated[
121
111
  str | None,
122
112
  typer.Option(
123
- help="Dataset split (e.g. train, validation). Auto-detected if omitted."
113
+ help=(
114
+ "Dataset split (e.g. train, validation). "
115
+ "If omitted, downloads all available splits."
116
+ )
124
117
  ),
125
118
  ] = None,
126
119
  ):
127
120
  if download_type == DownloadType.annotation:
128
- get_client(api_key).dataset_download(str(pk), export_type, split_type)
121
+ get_client(api_key).dataset_download(str(pk), split_type=split_type)
129
122
  else:
130
123
  typer.echo(get_client(api_key).dataset_download_files(str(pk)))
131
124
 
@@ -6,7 +6,7 @@ import logging
6
6
  import os
7
7
  import uuid
8
8
  from pathlib import Path
9
- from typing import Any
9
+ from typing import Any, ClassVar
10
10
  from urllib.parse import urlparse
11
11
 
12
12
  import httpx
@@ -197,44 +197,103 @@ class UnitlabClient:
197
197
  def project(self, project_id: str, pretty: int = 0) -> Any:
198
198
  return self._get(f"/api/sdk/projects/{project_id}/?pretty={pretty}")
199
199
 
200
- def project_members(self, project_id: str, pretty: int = 0) -> Any:
201
- return self._get(f"/api/sdk/projects/{project_id}/members/?pretty={pretty}")
202
-
203
200
  def project_upload_info(self, project_id: str) -> Any:
204
201
  return self._get(f"/api/sdk/projects/{project_id}/upload-info/")
205
202
 
203
+ # Supported upload extension groups. Mixed projects need per-file type
204
+ # detection so the SDK can attach only the options relevant to each file.
205
+ _EXTENSIONS_BY_GENERIC_TYPE: ClassVar[dict[str, set[str]]] = {
206
+ "img": {"jpg", "png", "jpeg", "webp", "gif", "bmp", "ico", "svg"},
207
+ "text": {"txt"},
208
+ "video": {"mp4", "avi", "mov", "webm", "mkv", "m4v", "wmv", "flv"},
209
+ "audio": {"mp3", "wav", "ogg", "aac", "flac", "m4a"},
210
+ "medical": {"dcm", "nii", "nii.gz", "nrrd"},
211
+ "document": {"pdf"},
212
+ }
213
+
214
+ @classmethod
215
+ def _extension_for_filename(cls, filename: str) -> str:
216
+ name = filename.lower()
217
+ if name.endswith(".nii.gz"):
218
+ return "nii.gz"
219
+ return name.rsplit(".", 1)[-1] if "." in name else ""
220
+
221
+ @classmethod
222
+ def _detect_generic_type(cls, filename: str) -> str | None:
223
+ """Resolve generic_type from a filename's extension.
224
+
225
+ Handles the compound suffix `nii.gz` by checking the last two dotted
226
+ components first. Returns None for unknown extensions; callers should
227
+ fall back to the project-level type or reject the file.
228
+ """
229
+ ext = cls._extension_for_filename(filename)
230
+ for generic_type, exts in cls._EXTENSIONS_BY_GENERIC_TYPE.items():
231
+ if ext in exts:
232
+ return generic_type
233
+ return None
234
+
235
+ @classmethod
236
+ def _known_file_extensions(cls) -> set[str]:
237
+ return {
238
+ extension
239
+ for extensions in cls._EXTENSIONS_BY_GENERIC_TYPE.values()
240
+ for extension in extensions
241
+ }
242
+
206
243
  def project_upload_data(
207
244
  self,
208
245
  project_id: str,
209
246
  directory: str | Path,
210
247
  batch_size: int = 100,
211
- sentences_per_chunk: int = 10,
212
- fps: float = 1.0,
213
- ) -> None:
248
+ fps: float = 1.0, # Applies only to video files.
249
+ ) -> dict[str, int]:
214
250
  directory = str(directory)
215
251
  if not os.path.isdir(directory):
216
252
  raise ValueError(f"Directory {directory} does not exist")
217
253
 
218
254
  upload_info = self.project_upload_info(project_id)
219
- accepted_formats: set[str] = set(upload_info["accepted_formats"])
220
- max_file_size_bytes: int = upload_info["max_file_size"]
221
- generic_type: str = upload_info["generic_type"]
255
+ accepted_formats: set[str] = {
256
+ str(ext).lower().lstrip(".") for ext in upload_info["accepted_formats"]
257
+ }
258
+ if not accepted_formats:
259
+ accepted_formats = self._known_file_extensions()
260
+ # Some projects do not declare a single type, so infer per file and
261
+ # use this only as a fallback for unrecognized accepted extensions.
262
+ project_generic_type: str | None = upload_info.get("generic_type")
263
+ max_file_sizes: dict[str, int] = {
264
+ str(generic_type): int(max_size)
265
+ for generic_type, max_size in (
266
+ upload_info.get("max_file_sizes") or {}
267
+ ).items()
268
+ }
269
+ fallback_max_file_size = upload_info.get("max_file_size")
270
+
271
+ def max_file_size_for(filename: str) -> int | None:
272
+ file_generic_type = (
273
+ self._detect_generic_type(filename) or project_generic_type
274
+ )
275
+ if file_generic_type and file_generic_type in max_file_sizes:
276
+ return max_file_sizes[file_generic_type]
277
+ if project_generic_type and fallback_max_file_size is not None:
278
+ return int(fallback_max_file_size)
279
+ return None
222
280
 
223
- # Single-pass file discovery with os.scandir
224
281
  files: list[str] = []
225
- with os.scandir(directory) as entries:
226
- for entry in entries:
227
- if not entry.is_file():
228
- continue
229
- ext = entry.name.rsplit(".", 1)[-1] if "." in entry.name else ""
230
- if ext.lower() in accepted_formats:
231
- files.append(entry.path)
282
+ for root, _dirs, filenames in os.walk(directory):
283
+ for filename in sorted(filenames):
284
+ ext = self._extension_for_filename(filename)
285
+ if ext in accepted_formats:
286
+ files.append(os.path.join(root, filename))
232
287
 
233
288
  filtered_files: list[str] = []
234
- max_file_size_mb = max_file_size_bytes / 1024 / 1024
235
289
  for file in files:
236
290
  file_size_bytes = os.path.getsize(file)
237
- if file_size_bytes > max_file_size_bytes:
291
+ max_file_size_bytes = max_file_size_for(os.path.basename(file))
292
+ if (
293
+ max_file_size_bytes is not None
294
+ and file_size_bytes > max_file_size_bytes
295
+ ):
296
+ max_file_size_mb = max_file_size_bytes / 1024 / 1024
238
297
  file_size_mb = file_size_bytes / 1024 / 1024
239
298
  logger.warning(
240
299
  f"File {file} is too large "
@@ -245,23 +304,40 @@ class UnitlabClient:
245
304
  filtered_files.append(file)
246
305
 
247
306
  num_files = len(filtered_files)
307
+ if num_files == 0:
308
+ raise ValueError(
309
+ "No uploadable files found. Accepted extensions: "
310
+ f"{', '.join(sorted(accepted_formats))}"
311
+ )
312
+
248
313
  num_batches = (num_files + batch_size - 1) // batch_size
249
314
  semaphore = asyncio.Semaphore(_UPLOAD_CONCURRENCY)
250
315
 
251
- # Medical projects use session-based upload with finalize
252
- session_id = str(uuid.uuid4()) if generic_type == "medical" else None
316
+ # Medical files in the same upload run share one session; non-medical
317
+ # files are sent without it.
318
+ medical_session_id = str(uuid.uuid4())
319
+
320
+ # Finalization is only relevant when this batch includes medical files.
321
+ has_medical_files = any(
322
+ self._detect_generic_type(os.path.basename(f)) == "medical"
323
+ for f in filtered_files
324
+ )
253
325
 
254
326
  async def post_file(
255
327
  client: httpx.AsyncClient, file: str, project_id: str
256
- ) -> int:
328
+ ) -> tuple[str, str | None, bool, str | None]:
257
329
  async with semaphore:
330
+ # Mixed uploads need per-file options; single-type projects
331
+ # can still fall back to the project-level type.
332
+ file_generic_type = (
333
+ self._detect_generic_type(os.path.basename(file))
334
+ or project_generic_type
335
+ )
258
336
  extra_data: dict[str, str] = {}
259
- if generic_type == "text":
260
- extra_data["sentences_per_chunk"] = str(sentences_per_chunk)
261
- elif generic_type == "video":
337
+ if file_generic_type == "video":
262
338
  extra_data["fps"] = str(fps)
263
- if session_id:
264
- extra_data["session_id"] = session_id
339
+ elif file_generic_type == "medical":
340
+ extra_data["session_id"] = medical_session_id
265
341
 
266
342
  try:
267
343
  with open(file, "rb") as f:
@@ -271,13 +347,22 @@ class UnitlabClient:
271
347
  data=extra_data,
272
348
  )
273
349
  response.raise_for_status()
274
- return 1
350
+ return file, file_generic_type, True, None
351
+ except httpx.HTTPStatusError as e:
352
+ return (
353
+ file,
354
+ file_generic_type,
355
+ False,
356
+ _extract_error_message(e.response),
357
+ )
275
358
  except Exception as e:
276
- logger.error(f"Error uploading file {file} - {e}")
277
- return 0
359
+ return file, file_generic_type, False, str(e)
278
360
 
279
- async def main() -> None:
361
+ async def main() -> dict[str, int]:
280
362
  logger.info(f"Uploading {num_files} files to project {project_id}")
363
+ uploaded = 0
364
+ uploaded_medical = 0
365
+ failed_files: list[tuple[str, str]] = []
281
366
  with tqdm.tqdm(total=num_files, ncols=80) as pbar:
282
367
  async with httpx.AsyncClient(
283
368
  base_url=self.api_url,
@@ -292,36 +377,80 @@ class UnitlabClient:
292
377
  ]
293
378
  ]
294
379
  for f in asyncio.as_completed(tasks):
295
- pbar.update(await f)
380
+ file, generic_type, success, error = await f
381
+ pbar.update(1)
382
+ if success:
383
+ uploaded += 1
384
+ if generic_type == "medical":
385
+ uploaded_medical += 1
386
+ else:
387
+ failed_files.append(
388
+ (file, error or "Unknown upload error")
389
+ )
390
+ logger.error(
391
+ "Error uploading file %s - %s",
392
+ file,
393
+ error or "Unknown upload error",
394
+ )
395
+
396
+ # Finalize only after at least one medical file succeeded.
397
+ if has_medical_files and uploaded_medical:
398
+ try:
399
+ response = await client.post(
400
+ f"/api/sdk/projects/{project_id}/"
401
+ f"medical-upload-sessions/{medical_session_id}/finalize/",
402
+ )
403
+ response.raise_for_status()
404
+ logger.info("Medical session finalized")
405
+ except httpx.HTTPStatusError as e:
406
+ raise NetworkError(
407
+ message=_extract_error_message(e.response),
408
+ detail=e,
409
+ ) from e
410
+ except httpx.HTTPError as e:
411
+ raise NetworkError(
412
+ message=f"Failed to finalize medical session: {e}",
413
+ detail=e,
414
+ ) from e
415
+
416
+ if failed_files:
417
+ first_failures = "; ".join(
418
+ f"{os.path.basename(path)}: {error}"
419
+ for path, error in failed_files[:5]
420
+ )
421
+ if len(failed_files) > 5:
422
+ first_failures += f"; and {len(failed_files) - 5} more"
423
+ raise NetworkError(
424
+ message=(
425
+ f"Failed to upload {len(failed_files)} of {num_files} files. "
426
+ f"{first_failures}"
427
+ )
428
+ )
296
429
 
297
- # Finalize medical session after all uploads
298
- if session_id:
299
- response = await client.post(
300
- f"/api/sdk/projects/{project_id}/"
301
- f"medical-upload-sessions/{session_id}/finalize/",
302
- )
303
- response.raise_for_status()
304
- logger.info("Medical session finalized")
430
+ return {
431
+ "total": num_files,
432
+ "uploaded": uploaded,
433
+ "failed": len(failed_files),
434
+ }
305
435
 
306
- asyncio.run(main())
436
+ return asyncio.run(main())
307
437
 
308
438
  def datasets(self, pretty: int = 0) -> Any:
309
- return self._get(f"/api/sdk/datasets/?pretty={pretty}")
439
+ return self._get(f"/api/sdk/releases/?pretty={pretty}")
310
440
 
311
441
  def dataset_download(
312
442
  self,
313
443
  dataset_id: str,
314
- export_type: str,
444
+ *,
315
445
  split_type: str | None = None,
316
446
  ) -> str:
317
447
  data: dict[str, Any] = {
318
448
  "download_type": "annotation",
319
- "export_type": export_type,
320
449
  }
321
450
  if split_type is not None:
322
451
  data["split_type"] = split_type
323
452
 
324
- response = self._post(f"/api/sdk/datasets/{dataset_id}/", data=data)
453
+ response = self._post(f"/api/sdk/releases/{dataset_id}/", data=data)
325
454
 
326
455
  file_url = response["file"]
327
456
  try:
@@ -349,7 +478,7 @@ class UnitlabClient:
349
478
  Path to the download folder.
350
479
  """
351
480
  response = self._post(
352
- f"/api/sdk/datasets/{dataset_id}/", data={"download_type": "files"}
481
+ f"/api/sdk/releases/{dataset_id}/", data={"download_type": "files"}
353
482
  )
354
483
  base_folder = str(dataset_id)
355
484
  os.makedirs(base_folder, exist_ok=True)
@@ -383,7 +512,7 @@ class UnitlabClient:
383
512
 
384
513
  async def download_file(
385
514
  client: httpx.AsyncClient, dataset_file: dict[str, str]
386
- ) -> int:
515
+ ) -> tuple[str, bool, str | None]:
387
516
  async with semaphore:
388
517
  try:
389
518
  async with client.stream("GET", dataset_file["source"]) as r:
@@ -391,14 +520,15 @@ class UnitlabClient:
391
520
  with open(dataset_file["file_path"], "wb") as f:
392
521
  async for chunk in r.aiter_bytes():
393
522
  f.write(chunk)
394
- return 1
523
+ return dataset_file["file_name"], True, None
395
524
  except Exception as e:
396
525
  logger.error(
397
526
  f"Error downloading file {dataset_file['file_name']} - {e}"
398
527
  )
399
- return 0
528
+ return dataset_file["file_name"], False, str(e)
400
529
 
401
530
  async def main() -> None:
531
+ failed_files: list[tuple[str, str]] = []
402
532
  with tqdm.tqdm(total=len(files_to_download), ncols=80) as pbar:
403
533
  async with httpx.AsyncClient(
404
534
  timeout=600.0, follow_redirects=True
@@ -408,7 +538,25 @@ class UnitlabClient:
408
538
  for df in files_to_download
409
539
  ]
410
540
  for f in asyncio.as_completed(tasks):
411
- pbar.update(await f)
541
+ file_name, success, error = await f
542
+ pbar.update(1)
543
+ if not success:
544
+ failed_files.append(
545
+ (file_name, error or "Unknown download error")
546
+ )
547
+
548
+ if failed_files:
549
+ first_failures = "; ".join(
550
+ f"{name}: {error}" for name, error in failed_files[:5]
551
+ )
552
+ if len(failed_files) > 5:
553
+ first_failures += f"; and {len(failed_files) - 5} more"
554
+ raise NetworkError(
555
+ message=(
556
+ f"Failed to download {len(failed_files)} of "
557
+ f"{len(files_to_download)} files. {first_failures}"
558
+ )
559
+ )
412
560
 
413
561
  asyncio.run(main())
414
562
  return base_folder
File without changes
File without changes
File without changes
File without changes
File without changes