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