download_cos.py
4.09 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
#!/usr/bin/env python3
import argparse
import csv
import os
import re
import sys
from pathlib import Path
from urllib.parse import urlparse
from dotenv import load_dotenv
from qcloud_cos import CosConfig, CosS3Client
load_dotenv(Path(__file__).resolve().parents[1] / '.env')
AUDIO_EXTS = {'.mp3', '.wav', '.flac', '.m4a', '.aac', '.ogg', '.wma', '.ape', '.alac'}
def normalize_key(raw_url: str) -> str:
if not raw_url:
return ''
raw_url = raw_url.strip()
if re.match(r'^https?://', raw_url, re.I):
return urlparse(raw_url).path.lstrip('/')
return raw_url.lstrip('/')
def ext_ok(path: str) -> bool:
return Path(path).suffix.lower() in AUDIO_EXTS
def main():
ap = argparse.ArgumentParser()
ap.add_argument('--manifest', default='output_selection_budgeted/selected_files.csv')
ap.add_argument('--output-dir', default='downloads')
ap.add_argument('--types', default='1,7,8,11,16')
ap.add_argument('--song-limit', type=int, default=3)
ap.add_argument('--start-song-offset', type=int, default=0)
ap.add_argument('--overwrite', action='store_true')
ap.add_argument('--fail-log', default='download_failures.csv')
args = ap.parse_args()
region = os.environ['COS_REGION']
secret_id = os.environ['COS_SECRET_ID']
secret_key = os.environ['COS_SECRET_KEY']
bucket = os.environ['COS_BUCKET']
config = CosConfig(Region=region, SecretId=secret_id, SecretKey=secret_key)
client = CosS3Client(config)
wanted_types = {int(x) for x in args.types.split(',') if x.strip()}
out_dir = Path(args.output_dir)
out_dir.mkdir(parents=True, exist_ok=True)
fail_log = Path(args.fail_log)
fail_exists = fail_log.exists()
fail_f = fail_log.open('a', encoding='utf-8', newline='')
fail_w = csv.writer(fail_f)
if not fail_exists:
fail_w.writerow(['song_id', 'type', 'url', 'error'])
downloaded = skipped = errors = 0
active_song_ids = []
current_song = None
started = False
song_counter = 0
with open(args.manifest, 'r', encoding='utf-8', newline='') as f:
reader = csv.DictReader(f)
for row in reader:
sid = row['song_id']
if sid != current_song:
current_song = sid
if song_counter < args.start_song_offset:
song_counter += 1
continue
if args.song_limit and started and len(active_song_ids) >= args.song_limit and sid not in set(active_song_ids):
break
if sid not in active_song_ids:
active_song_ids.append(sid)
started = True
if sid not in active_song_ids:
continue
try:
t = int(row['type'])
except Exception:
skipped += 1
continue
if t not in wanted_types:
skipped += 1
continue
key = normalize_key(row['url'])
if not key or not ext_ok(key):
skipped += 1
continue
target_dir = out_dir / sid / f'type_{t}'
target_dir.mkdir(parents=True, exist_ok=True)
target = target_dir / Path(key).name
if target.exists() and not args.overwrite:
skipped += 1
continue
try:
client.download_file(Bucket=bucket, Key=key, DestFilePath=str(target))
downloaded += 1
if downloaded % 50 == 0:
print({'downloaded': downloaded, 'songs': len(active_song_ids), 'last': str(target)})
except Exception as e:
errors += 1
fail_w.writerow([sid, t, row['url'], str(e)])
fail_f.flush()
print({'error_song': sid, 'type': t, 'key': key, 'error': str(e)}, file=sys.stderr)
fail_f.close()
print({
'downloaded_files': downloaded,
'skipped_rows': skipped,
'error_files': errors,
'song_count': len(active_song_ids),
'output_dir': str(out_dir.resolve()),
})
if __name__ == '__main__':
main()