2Utility functions for downloading from S3 public buckets
3"""
4import boto3
5from botocore import UNSIGNED
6from botocore.config import Config
7import os
8from pathlib import Path
9import sys
12class ProgressCallback:
13 """Callback to show download progress"""
14 def __init__(self, filename, filesize):
15 self._filename = filename
16 self._size = filesize
17 self._seen_so_far = 0
19 def __call__(self, bytes_amount):
20 self._seen_so_far += bytes_amount
21 percentage = (self._seen_so_far / self._size) * 100 if self._size > 0 else 0
22 sys.stdout.write(
23 f"\r Progress: {self._seen_so_far:,} / {self._size:,} bytes ({percentage:.1f}%)"
24 )
25 sys.stdout.flush()
28def download_s3_folder(s3_url: str, local_dir: str, skip_confirmation: bool = False):
29 """
30 Download entire folder from S3 public bucket
32 Args:
33 s3_url: S3 URL in format s3://bucket-name/path/to/folder/
34 local_dir: Local directory path to download files to
35 skip_confirmation: If True, skip user confirmation prompt
36 """
37 # Parse S3 URL
38 if not s3_url.startswith('s3://'):
39 raise ValueError(f"S3 URL must start with 's3://': {s3_url}")
41 s3_url = s3_url.rstrip('/')
42 s3_parts = s3_url[5:].split('/', 1)
43 bucket_name = s3_parts[0]
44 prefix = s3_parts[1] + '/' if len(s3_parts) > 1 else ''
46 # S3 configuration for public bucket (no credentials needed)
47 s3 = boto3.client('s3', config=Config(signature_version=UNSIGNED))
49 # Create local directory
50 Path(local_dir).mkdir(parents=True, exist_ok=True)
52 print(f"Scanning s3://{bucket_name}/{prefix}")
53 print(f"Will download to local directory: {local_dir}/")
54 print()
56 # List all objects in the folder to calculate total size
57 paginator = s3.get_paginator('list_objects_v2')
58 pages = paginator.paginate(Bucket=bucket_name, Prefix=prefix)
60 files_to_download = []
61 total_size = 0
63 for page in pages:
64 if 'Contents' not in page:
65 print("No files found in the specified path")
66 return
68 for obj in page['Contents']:
69 s3_key = obj['Key']
70 file_size = obj['Size']
72 # Skip if it's just the directory itself
73 if s3_key == prefix or s3_key == prefix.rstrip('/'):
74 continue
76 # Get the relative path (remove the prefix)
77 relative_path = s3_key[len(prefix):]
78 if not relative_path:
79 continue
81 files_to_download.append({
82 's3_key': s3_key,
83 'relative_path': relative_path,
84 'size': file_size
85 })
86 total_size += file_size
88 print(f"Found {len(files_to_download)} files")
89 print(f"Total size: {total_size:,} bytes ({total_size / (1024**2):.2f} MB, {total_size / (1024**3):.2f} GB)")
90 print()
92 if not skip_confirmation:
93 response = input("Continue with download? (y/n): ")
94 if response.lower() != 'y':
95 print("Download cancelled")
96 return
98 print()
99 print("Starting download...")
100 print()
102 # Download all files
103 downloaded_bytes = 0
104 for idx, file_info in enumerate(files_to_download, 1):
105 s3_key = file_info['s3_key']
106 relative_path = file_info['relative_path']
107 file_size = file_info['size']
109 # Local file path
110 local_file = os.path.join(local_dir, relative_path)
112 # Create subdirectories if needed
113 local_file_dir = os.path.dirname(local_file)
114 if local_file_dir:
115 Path(local_file_dir).mkdir(parents=True, exist_ok=True)
117 # Download the file with progress callback
427c40eprepare datasetsJeremy Magland 118 # Using download_fileobj to write directly to file for partial download support
6864793support multi-channelJeremy Magland 119 print(f"[{idx}/{len(files_to_download)}] {relative_path} ({file_size:,} bytes)")
120 progress = ProgressCallback(relative_path, file_size)
122 s3.download_fileobj(bucket_name, s3_key, f, Callback=progress)
125 downloaded_bytes += file_size
126 overall_progress = (downloaded_bytes / total_size) * 100
127 print(f" Overall progress: {downloaded_bytes:,} / {total_size:,} bytes ({overall_progress:.1f}%)")
128 print()
130 print(f"Download complete!")
131 print(f"Total files: {len(files_to_download)}")
132 print(f"Total size: {total_size:,} bytes ({total_size / (1024**3):.2f} GB)")