DiscogsDataProcessorCLI 1.5.0__py3-none-any.whl
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.
- discogs/__init__.py +0 -0
- discogs/chunker.py +95 -0
- discogs/config.py +75 -0
- discogs/converter.py +190 -0
- discogs/deleter.py +51 -0
- discogs/downloader.py +120 -0
- discogs/extractor.py +92 -0
- discogs/main.py +170 -0
- discogs/scraper.py +125 -0
- discogs/selector.py +197 -0
- discogs/utils.py +96 -0
- discogsdataprocessorcli-1.5.0.dist-info/METADATA +119 -0
- discogsdataprocessorcli-1.5.0.dist-info/RECORD +16 -0
- discogsdataprocessorcli-1.5.0.dist-info/WHEEL +5 -0
- discogsdataprocessorcli-1.5.0.dist-info/entry_points.txt +2 -0
- discogsdataprocessorcli-1.5.0.dist-info/top_level.txt +1 -0
discogs/__init__.py
ADDED
|
File without changes
|
discogs/chunker.py
ADDED
|
@@ -0,0 +1,95 @@
|
|
|
1
|
+
# discogs/chunker.py
|
|
2
|
+
|
|
3
|
+
import re
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
from rich.console import Console
|
|
6
|
+
from rich.progress import Progress, BarColumn, TimeElapsedColumn, TextColumn
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
def sanitize_line(line: str) -> str:
|
|
10
|
+
"""
|
|
11
|
+
Removes invalid XML characters and fixes unescaped ampersands.
|
|
12
|
+
This helps ensure the XML is well-formed before processing.
|
|
13
|
+
"""
|
|
14
|
+
line = re.sub(r'[^\x09\x0A\x0D\x20-\uD7FF\uE000-\uFFFD]', '', line) # Remove illegal XML characters
|
|
15
|
+
line = re.sub(r'&(?![a-zA-Z0-9#]+;)', '&', line) # Escape unescaped '&' characters
|
|
16
|
+
return line
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
def chunk_xml_by_type(xml_file: Path, content_type: str, records_per_file: int = 10000) -> Path:
|
|
20
|
+
"""
|
|
21
|
+
Splits a large XML file into smaller, valid XML files (chunks).
|
|
22
|
+
Each chunk contains up to `records_per_file` XML records.
|
|
23
|
+
Returns the folder path where chunked files are stored.
|
|
24
|
+
"""
|
|
25
|
+
record_tag = content_type[:-1].lower() # e.g., "releases" → "release"
|
|
26
|
+
start_pat = re.compile(fr'<{record_tag}\b', re.IGNORECASE) # Match opening tag
|
|
27
|
+
end_pat = re.compile(fr'</{record_tag}>', re.IGNORECASE) # Match closing tag
|
|
28
|
+
|
|
29
|
+
chunk_folder = xml_file.parent / f"chunked_{content_type}" # Output folder
|
|
30
|
+
chunk_folder.mkdir(parents=True, exist_ok=True)
|
|
31
|
+
|
|
32
|
+
console = Console()
|
|
33
|
+
chunk_count = 0
|
|
34
|
+
record_count = 0
|
|
35
|
+
inside_record = False
|
|
36
|
+
buffer_lines = [] # Stores lines of current XML record
|
|
37
|
+
current_chunk_file = None
|
|
38
|
+
|
|
39
|
+
# Helper function to open a new chunk file
|
|
40
|
+
def open_new_chunk():
|
|
41
|
+
nonlocal chunk_count, current_chunk_file, record_count
|
|
42
|
+
chunk_count += 1
|
|
43
|
+
chunk_path = chunk_folder / f"chunk_{chunk_count:05}.xml"
|
|
44
|
+
current_chunk_file = open(chunk_path, "w", encoding="utf-8")
|
|
45
|
+
current_chunk_file.write(f'<?xml version="1.0" encoding="utf-8"?>\n<{content_type}>\n')
|
|
46
|
+
record_count = 0
|
|
47
|
+
|
|
48
|
+
# Helper function to close the current chunk file
|
|
49
|
+
def close_chunk():
|
|
50
|
+
nonlocal current_chunk_file
|
|
51
|
+
if current_chunk_file:
|
|
52
|
+
current_chunk_file.write(f"</{content_type}>")
|
|
53
|
+
current_chunk_file.close()
|
|
54
|
+
current_chunk_file = None
|
|
55
|
+
|
|
56
|
+
open_new_chunk()
|
|
57
|
+
|
|
58
|
+
# Setup progress bar for visual feedback
|
|
59
|
+
with Progress(
|
|
60
|
+
TextColumn("[progress.description]{task.description}"),
|
|
61
|
+
BarColumn(),
|
|
62
|
+
"[progress.percentage]{task.percentage:.1f}%",
|
|
63
|
+
"•",
|
|
64
|
+
TimeElapsedColumn()
|
|
65
|
+
) as progress:
|
|
66
|
+
task = progress.add_task(f"Chunking {xml_file.name}", total=xml_file.stat().st_size)
|
|
67
|
+
|
|
68
|
+
# Read and process XML file line by line
|
|
69
|
+
with xml_file.open("r", encoding="utf-8", errors="ignore") as f:
|
|
70
|
+
for raw_line in f:
|
|
71
|
+
line = sanitize_line(raw_line)
|
|
72
|
+
progress.update(task, advance=len(raw_line))
|
|
73
|
+
|
|
74
|
+
if not inside_record:
|
|
75
|
+
# Detect start of a record
|
|
76
|
+
if start_pat.search(line):
|
|
77
|
+
inside_record = True
|
|
78
|
+
buffer_lines = [line]
|
|
79
|
+
else:
|
|
80
|
+
buffer_lines.append(line)
|
|
81
|
+
if end_pat.search(line):
|
|
82
|
+
# Write complete record to current chunk
|
|
83
|
+
current_chunk_file.write("".join(buffer_lines) + "\n")
|
|
84
|
+
record_count += 1
|
|
85
|
+
inside_record = False
|
|
86
|
+
buffer_lines = []
|
|
87
|
+
|
|
88
|
+
# If chunk is full, start a new one
|
|
89
|
+
if record_count >= records_per_file:
|
|
90
|
+
close_chunk()
|
|
91
|
+
open_new_chunk()
|
|
92
|
+
|
|
93
|
+
close_chunk()
|
|
94
|
+
console.print(f"[green]✔ Chunked into {chunk_count} file(s): {chunk_folder}")
|
|
95
|
+
return chunk_folder
|
discogs/config.py
ADDED
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
# discogs/config.py
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
from rich.console import Console
|
|
6
|
+
from rich.prompt import Prompt
|
|
7
|
+
|
|
8
|
+
CONFIG_PATH = Path.home() / ".discogs_config.json" # Path to the user's config file
|
|
9
|
+
DEFAULT_DOWNLOAD_PATH = Path.home() / "Downloads" / "Discogs" # Default download location
|
|
10
|
+
|
|
11
|
+
console = Console()
|
|
12
|
+
|
|
13
|
+
def get_download_dir() -> Path:
|
|
14
|
+
"""
|
|
15
|
+
Returns the download directory from the config file if it exists,
|
|
16
|
+
otherwise returns the default download path.
|
|
17
|
+
"""
|
|
18
|
+
if CONFIG_PATH.exists():
|
|
19
|
+
path = CONFIG_PATH.read_text().strip()
|
|
20
|
+
if path:
|
|
21
|
+
return Path(path)
|
|
22
|
+
return DEFAULT_DOWNLOAD_PATH
|
|
23
|
+
|
|
24
|
+
def set_download_dir(path: str) -> None:
|
|
25
|
+
"""
|
|
26
|
+
Sets (writes) the download directory path to the config file.
|
|
27
|
+
"""
|
|
28
|
+
CONFIG_PATH.write_text(path.strip())
|
|
29
|
+
|
|
30
|
+
def load_config() -> dict:
|
|
31
|
+
"""
|
|
32
|
+
Loads configuration data from the config file.
|
|
33
|
+
If the config cannot be read, it returns a dictionary with default values.
|
|
34
|
+
"""
|
|
35
|
+
if CONFIG_PATH.exists():
|
|
36
|
+
try:
|
|
37
|
+
with CONFIG_PATH.open("r") as f:
|
|
38
|
+
return json.load(f)
|
|
39
|
+
except Exception as e:
|
|
40
|
+
console.print(f"[red]Error reading config:[/] {e}")
|
|
41
|
+
return {"download_dir": str(DEFAULT_DOWNLOAD_PATH)}
|
|
42
|
+
|
|
43
|
+
def save_config(config: dict) -> None:
|
|
44
|
+
"""
|
|
45
|
+
Saves the given configuration dictionary to the config file.
|
|
46
|
+
Creates the parent directory if it doesn't exist.
|
|
47
|
+
"""
|
|
48
|
+
CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
49
|
+
with CONFIG_PATH.open("w") as f:
|
|
50
|
+
json.dump(config, f, indent=4)
|
|
51
|
+
|
|
52
|
+
def get_download_dir() -> Path:
|
|
53
|
+
"""
|
|
54
|
+
Gets the download directory from config, falling back to default if missing.
|
|
55
|
+
"""
|
|
56
|
+
config = load_config()
|
|
57
|
+
return Path(config.get("download_dir", str(DEFAULT_DOWNLOAD_PATH)))
|
|
58
|
+
|
|
59
|
+
def configure_download_folder() -> None:
|
|
60
|
+
"""
|
|
61
|
+
Prompts the user to configure the download folder via CLI.
|
|
62
|
+
Updates the config if a new folder is provided.
|
|
63
|
+
"""
|
|
64
|
+
console.print("[bold]Configure Discogs download folder[/bold]")
|
|
65
|
+
current = get_download_dir()
|
|
66
|
+
console.print(f"Current folder: [green]{current}[/green]")
|
|
67
|
+
|
|
68
|
+
new_path = Prompt.ask("Enter new folder path or leave empty to keep current", default=str(current)).strip()
|
|
69
|
+
if new_path:
|
|
70
|
+
new_path = Path(new_path).expanduser()
|
|
71
|
+
new_path.mkdir(parents=True, exist_ok=True)
|
|
72
|
+
save_config({"download_dir": str(new_path)})
|
|
73
|
+
console.print(f"[green]✅ Download folder updated to:[/] {new_path}")
|
|
74
|
+
else:
|
|
75
|
+
console.print("[yellow]No changes made.[/yellow]")
|
discogs/converter.py
ADDED
|
@@ -0,0 +1,190 @@
|
|
|
1
|
+
# discogs/converter.py
|
|
2
|
+
|
|
3
|
+
import shutil
|
|
4
|
+
import json
|
|
5
|
+
import csv
|
|
6
|
+
import xml.etree.ElementTree as ET
|
|
7
|
+
from pathlib import Path
|
|
8
|
+
from rich.console import Console
|
|
9
|
+
from rich.progress import (
|
|
10
|
+
Progress,
|
|
11
|
+
SpinnerColumn,
|
|
12
|
+
BarColumn,
|
|
13
|
+
TextColumn,
|
|
14
|
+
TimeElapsedColumn,
|
|
15
|
+
DownloadColumn,
|
|
16
|
+
TransferSpeedColumn,
|
|
17
|
+
)
|
|
18
|
+
|
|
19
|
+
from discogs.chunker import chunk_xml_by_type
|
|
20
|
+
|
|
21
|
+
console = Console()
|
|
22
|
+
|
|
23
|
+
def _scan_columns(chunk_file: Path, record_tag: str, column_set: set):
|
|
24
|
+
"""
|
|
25
|
+
Scans an XML chunk file to identify all unique tag paths and attributes.
|
|
26
|
+
Adds these as potential CSV columns.
|
|
27
|
+
"""
|
|
28
|
+
current_path = []
|
|
29
|
+
for event, elem in ET.iterparse(chunk_file, events=("start", "end")):
|
|
30
|
+
if event == "start":
|
|
31
|
+
current_path.append(elem.tag)
|
|
32
|
+
# Add all attributes of the current tag to the column set
|
|
33
|
+
for attr in elem.attrib:
|
|
34
|
+
key = "_".join(current_path[-2:] + [attr]) if len(current_path) >= 2 else f"{elem.tag}_{attr}"
|
|
35
|
+
column_set.add(key)
|
|
36
|
+
elif event == "end":
|
|
37
|
+
if elem.text and not elem.text.isspace():
|
|
38
|
+
# Add tag path for text content
|
|
39
|
+
key = "_".join(current_path[-2:] + [elem.tag]) if len(current_path) >= 2 else elem.tag
|
|
40
|
+
column_set.add(key)
|
|
41
|
+
current_path.pop()
|
|
42
|
+
elem.clear()
|
|
43
|
+
|
|
44
|
+
def _write_rows(chunk_file: Path, writer: csv.DictWriter, columns: list, record_tag: str):
|
|
45
|
+
"""
|
|
46
|
+
Parses an XML chunk and writes each record as a CSV row using the given column list.
|
|
47
|
+
"""
|
|
48
|
+
current_path = []
|
|
49
|
+
record_data = {}
|
|
50
|
+
nested = {}
|
|
51
|
+
|
|
52
|
+
for event, elem in ET.iterparse(chunk_file, events=("start", "end")):
|
|
53
|
+
if event == "start":
|
|
54
|
+
current_path.append(elem.tag)
|
|
55
|
+
# Collect attribute values
|
|
56
|
+
for attr, val in elem.attrib.items():
|
|
57
|
+
key = "_".join(current_path[-2:] + [attr]) if len(current_path) >= 2 else f"{elem.tag}_{attr}"
|
|
58
|
+
nested.setdefault(key, []).append(val)
|
|
59
|
+
elif event == "end":
|
|
60
|
+
if elem.text and not elem.text.isspace():
|
|
61
|
+
# Collect text values
|
|
62
|
+
key = "_".join(current_path[-2:] + [elem.tag]) if len(current_path) >= 2 else elem.tag
|
|
63
|
+
nested.setdefault(key, []).append(elem.text.strip())
|
|
64
|
+
|
|
65
|
+
# End of a full record → flush it to CSV
|
|
66
|
+
if elem.tag == record_tag:
|
|
67
|
+
for k, v in nested.items():
|
|
68
|
+
record_data[k] = v[0] if len(v) == 1 else json.dumps(v) # Use first or serialize list
|
|
69
|
+
writer.writerow({col: record_data.get(col, "") for col in columns})
|
|
70
|
+
record_data.clear()
|
|
71
|
+
nested.clear()
|
|
72
|
+
|
|
73
|
+
current_path.pop()
|
|
74
|
+
elem.clear()
|
|
75
|
+
|
|
76
|
+
from time import perf_counter
|
|
77
|
+
|
|
78
|
+
def convert_chunks_to_csv(chunk_dir: Path, output_csv: Path, content_type: str):
|
|
79
|
+
"""
|
|
80
|
+
Converts all chunked XML files in a given folder into a single CSV file.
|
|
81
|
+
The function discovers all columns, parses each chunk, and writes rows.
|
|
82
|
+
"""
|
|
83
|
+
record_tag = content_type[:-1] # e.g. "releases" → "release"
|
|
84
|
+
chunks = sorted(chunk_dir.glob("chunk_*.xml"))
|
|
85
|
+
|
|
86
|
+
if not chunks:
|
|
87
|
+
console.print(f"[red]No XML chunks found in {chunk_dir}[/red]")
|
|
88
|
+
return
|
|
89
|
+
|
|
90
|
+
start_time = perf_counter()
|
|
91
|
+
|
|
92
|
+
# Step 1: Scan all chunks to detect all column names
|
|
93
|
+
column_set = set()
|
|
94
|
+
console.print("[bold]Step 1:[/] Scanning tags...")
|
|
95
|
+
|
|
96
|
+
with Progress(
|
|
97
|
+
SpinnerColumn(),
|
|
98
|
+
TextColumn("[progress.description]{task.description}"),
|
|
99
|
+
BarColumn(),
|
|
100
|
+
"[progress.percentage]{task.percentage:.1f}%",
|
|
101
|
+
"•",
|
|
102
|
+
TimeElapsedColumn()
|
|
103
|
+
) as p:
|
|
104
|
+
task = p.add_task("Scanning...", total=len(chunks))
|
|
105
|
+
for chunk in chunks:
|
|
106
|
+
_scan_columns(chunk, record_tag, column_set)
|
|
107
|
+
p.update(task, advance=1)
|
|
108
|
+
|
|
109
|
+
columns = sorted(column_set)
|
|
110
|
+
|
|
111
|
+
# Step 2: Write rows into CSV
|
|
112
|
+
console.print(f"[bold]Step 2:[/] Writing [green]{output_csv.name}[/green] with {len(columns)} columns...")
|
|
113
|
+
|
|
114
|
+
with open(output_csv, "w", newline="", encoding="utf-8") as f:
|
|
115
|
+
writer = csv.DictWriter(f, fieldnames=columns)
|
|
116
|
+
writer.writeheader()
|
|
117
|
+
|
|
118
|
+
with Progress(
|
|
119
|
+
SpinnerColumn(),
|
|
120
|
+
TextColumn("[progress.description]{task.description}"),
|
|
121
|
+
BarColumn(),
|
|
122
|
+
"[progress.percentage]{task.percentage:.1f}%",
|
|
123
|
+
"•",
|
|
124
|
+
TimeElapsedColumn()
|
|
125
|
+
) as p:
|
|
126
|
+
task = p.add_task("Converting...", total=len(chunks))
|
|
127
|
+
for chunk in chunks:
|
|
128
|
+
_write_rows(chunk, writer, columns, record_tag)
|
|
129
|
+
p.update(task, advance=1)
|
|
130
|
+
|
|
131
|
+
duration = perf_counter() - start_time
|
|
132
|
+
output_size_mb = output_csv.stat().st_size / (1024 * 1024)
|
|
133
|
+
|
|
134
|
+
# Final status output
|
|
135
|
+
console.print(f"\n[green]✔ CSV saved:[/] {output_csv}")
|
|
136
|
+
console.print("[bold green]✔ Conversion completed[/bold green]")
|
|
137
|
+
console.print(f"[bold white]📄 Chunks processed:[/] {len(chunks)} files")
|
|
138
|
+
console.print(f"[bold white]🧩 Output CSV:[/] {output_csv.name}")
|
|
139
|
+
console.print(f"[bold white]💾 Output size:[/] {output_size_mb:.2f} MB")
|
|
140
|
+
console.print(f"[bold white]🗂 Saved to:[/] {output_csv.parent}")
|
|
141
|
+
console.print(f"[bold white]⏱ Duration:[/] {duration:.1f} seconds")
|
|
142
|
+
|
|
143
|
+
def convert_xml_to_csv(xml_path: Path, content_type: str) -> Path:
|
|
144
|
+
"""
|
|
145
|
+
Full pipeline: chunk an XML file and convert the chunks to a CSV file.
|
|
146
|
+
Temporary chunked files are deleted after the process.
|
|
147
|
+
"""
|
|
148
|
+
chunk_dir = xml_path.parent / f"chunked_{content_type}"
|
|
149
|
+
output_csv = xml_path.with_suffix(".csv")
|
|
150
|
+
|
|
151
|
+
chunk_xml_by_type(xml_path, content_type) # Split large XML into smaller parts
|
|
152
|
+
convert_chunks_to_csv(chunk_dir, output_csv, content_type) # Convert chunks to CSV
|
|
153
|
+
shutil.rmtree(chunk_dir, ignore_errors=True) # Cleanup
|
|
154
|
+
|
|
155
|
+
return output_csv
|
|
156
|
+
|
|
157
|
+
def convert_interactively():
|
|
158
|
+
"""
|
|
159
|
+
Prompts user to select XML files for conversion.
|
|
160
|
+
"""
|
|
161
|
+
from rich.prompt import Prompt
|
|
162
|
+
from discogs.config import get_download_dir
|
|
163
|
+
from discogs.utils import open_folder
|
|
164
|
+
|
|
165
|
+
download_dir = get_download_dir()
|
|
166
|
+
dataset_dir = download_dir / "Datasets"
|
|
167
|
+
|
|
168
|
+
xml_files = list(dataset_dir.rglob("*.xml"))
|
|
169
|
+
if not xml_files:
|
|
170
|
+
console.print("[red]No XML files found to convert.[/red]")
|
|
171
|
+
return
|
|
172
|
+
|
|
173
|
+
console.print("[bold]Select XML file to convert:[/bold]")
|
|
174
|
+
for i, file in enumerate(xml_files):
|
|
175
|
+
console.print(f"[{i + 1}] {file.relative_to(download_dir)}")
|
|
176
|
+
|
|
177
|
+
choice = Prompt.ask("Enter number", default="1")
|
|
178
|
+
try:
|
|
179
|
+
idx = int(choice.strip()) - 1
|
|
180
|
+
if 0 <= idx < len(xml_files):
|
|
181
|
+
file = xml_files[idx]
|
|
182
|
+
content_type = file.stem.split("_")[-1]
|
|
183
|
+
convert_xml_to_csv(file, content_type)
|
|
184
|
+
open_folder(file.parent)
|
|
185
|
+
else:
|
|
186
|
+
console.print("[red]Invalid selection.[/red]")
|
|
187
|
+
except:
|
|
188
|
+
console.print("[red]Invalid input.[/red]")
|
|
189
|
+
|
|
190
|
+
__all__ = ["convert_xml_to_csv"] # Exported symbols
|
discogs/deleter.py
ADDED
|
@@ -0,0 +1,51 @@
|
|
|
1
|
+
# discogs/deleter.py
|
|
2
|
+
|
|
3
|
+
from pathlib import Path
|
|
4
|
+
from rich.console import Console
|
|
5
|
+
from discogs.selector import display_status_table, select_indices
|
|
6
|
+
from discogs.config import get_download_dir
|
|
7
|
+
from discogs.scraper import get_latest_files
|
|
8
|
+
|
|
9
|
+
console = Console()
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def delete_files():
|
|
13
|
+
"""
|
|
14
|
+
Interactive function that allows the user to delete selected
|
|
15
|
+
.gz, .xml, and .csv files from the dataset directory.
|
|
16
|
+
"""
|
|
17
|
+
download_dir = get_download_dir() # Get the base download directory from config
|
|
18
|
+
df = get_latest_files() # Load the latest file list (as a DataFrame)
|
|
19
|
+
|
|
20
|
+
if df.empty:
|
|
21
|
+
console.print("[red]No files found.[/red]")
|
|
22
|
+
return
|
|
23
|
+
|
|
24
|
+
# Show the current file status table
|
|
25
|
+
display_status_table(df, download_dir)
|
|
26
|
+
|
|
27
|
+
# Let the user select which files to delete
|
|
28
|
+
selected = select_indices(df)
|
|
29
|
+
|
|
30
|
+
for i in selected:
|
|
31
|
+
row = df.iloc[i]
|
|
32
|
+
filename = Path(row["url"]).name
|
|
33
|
+
year_month = row["month"]
|
|
34
|
+
data_dir = download_dir / "Datasets" / year_month
|
|
35
|
+
|
|
36
|
+
gz_path = data_dir / filename # Original .gz file
|
|
37
|
+
xml_path = gz_path.with_suffix("") # Extracted .xml file
|
|
38
|
+
csv_path = xml_path.with_suffix(".csv") # Converted .csv file
|
|
39
|
+
|
|
40
|
+
# Try deleting each file, one by one
|
|
41
|
+
for file in [gz_path, xml_path, csv_path]:
|
|
42
|
+
if file.exists():
|
|
43
|
+
file.unlink() # Delete the file
|
|
44
|
+
console.print(f"[red]🗑 Deleted:[/] {file.name}")
|
|
45
|
+
else:
|
|
46
|
+
console.print(f"[dim]• Not found:[/] {file.name}")
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
# Allow this script to be run directly
|
|
50
|
+
if __name__ == "__main__":
|
|
51
|
+
delete_files()
|
discogs/downloader.py
ADDED
|
@@ -0,0 +1,120 @@
|
|
|
1
|
+
# discogs/downloader.py
|
|
2
|
+
|
|
3
|
+
import time
|
|
4
|
+
import requests
|
|
5
|
+
from time import sleep
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
from urllib.parse import urlparse
|
|
8
|
+
from rich.console import Console
|
|
9
|
+
from rich.progress import (
|
|
10
|
+
Progress, BarColumn, DownloadColumn, TransferSpeedColumn,
|
|
11
|
+
TimeRemainingColumn, TextColumn, SpinnerColumn
|
|
12
|
+
)
|
|
13
|
+
from concurrent.futures import ThreadPoolExecutor, as_completed
|
|
14
|
+
from datetime import datetime
|
|
15
|
+
|
|
16
|
+
console = Console()
|
|
17
|
+
|
|
18
|
+
def _download_file(url: str, target_path: Path, progress, task_id, retries: int = 5) -> Path:
|
|
19
|
+
"""
|
|
20
|
+
Downloads a file with support for resume and retry.
|
|
21
|
+
Updates a Rich progress bar during download.
|
|
22
|
+
"""
|
|
23
|
+
headers = {}
|
|
24
|
+
downloaded = 0
|
|
25
|
+
|
|
26
|
+
# If file exists, resume from where it left off
|
|
27
|
+
if target_path.exists():
|
|
28
|
+
downloaded = target_path.stat().st_size
|
|
29
|
+
headers["Range"] = f"bytes={downloaded}-"
|
|
30
|
+
|
|
31
|
+
total_size = int(requests.head(url).headers.get("Content-Length", 0))
|
|
32
|
+
|
|
33
|
+
for attempt in range(retries):
|
|
34
|
+
try:
|
|
35
|
+
with requests.get(url, headers=headers, stream=True, timeout=10) as response:
|
|
36
|
+
response.raise_for_status()
|
|
37
|
+
|
|
38
|
+
mode = "ab" if downloaded else "wb"
|
|
39
|
+
with open(target_path, mode) as f:
|
|
40
|
+
for chunk in response.iter_content(chunk_size=1024 * 64):
|
|
41
|
+
if chunk:
|
|
42
|
+
f.write(chunk)
|
|
43
|
+
downloaded += len(chunk)
|
|
44
|
+
progress.update(task_id, completed=downloaded)
|
|
45
|
+
|
|
46
|
+
return target_path # Download completed
|
|
47
|
+
|
|
48
|
+
except requests.RequestException as e:
|
|
49
|
+
# Retry a few times if download fails
|
|
50
|
+
if attempt < retries - 1:
|
|
51
|
+
sleep(1.5)
|
|
52
|
+
continue
|
|
53
|
+
else:
|
|
54
|
+
raise RuntimeError(f"Download failed after {retries} retries: {e}")
|
|
55
|
+
|
|
56
|
+
def download_files_threaded(df, selected_indexes, download_dir: Path) -> list[Path]:
|
|
57
|
+
"""
|
|
58
|
+
Downloads multiple files concurrently using threads.
|
|
59
|
+
Displays a combined progress bar for all downloads.
|
|
60
|
+
"""
|
|
61
|
+
urls = [df.iloc[i]["url"] for i in selected_indexes]
|
|
62
|
+
paths = []
|
|
63
|
+
total_bytes = 0
|
|
64
|
+
start_time = time.time()
|
|
65
|
+
|
|
66
|
+
with Progress(
|
|
67
|
+
SpinnerColumn(),
|
|
68
|
+
TextColumn("[progress.description]{task.description} → [bold blue]{task.fields[filename]}", justify="right"),
|
|
69
|
+
BarColumn(),
|
|
70
|
+
"[progress.percentage]{task.percentage:>3.1f}%",
|
|
71
|
+
"•",
|
|
72
|
+
DownloadColumn(),
|
|
73
|
+
"•",
|
|
74
|
+
TransferSpeedColumn(),
|
|
75
|
+
"•",
|
|
76
|
+
TimeRemainingColumn(),
|
|
77
|
+
) as progress:
|
|
78
|
+
with ThreadPoolExecutor(max_workers=8) as executor:
|
|
79
|
+
futures = []
|
|
80
|
+
|
|
81
|
+
for url in urls:
|
|
82
|
+
filename = Path(urlparse(url).path).name
|
|
83
|
+
|
|
84
|
+
# Extract date from filename and create target folder
|
|
85
|
+
date_str = filename.split("_")[1]
|
|
86
|
+
year_month = datetime.strptime(date_str, "%Y%m%d").strftime("%Y-%m")
|
|
87
|
+
target_folder = download_dir / "Datasets" / year_month
|
|
88
|
+
target_folder.mkdir(parents=True, exist_ok=True)
|
|
89
|
+
target_path = target_folder / filename
|
|
90
|
+
|
|
91
|
+
# Skip already downloaded files
|
|
92
|
+
if target_path.exists():
|
|
93
|
+
console.print(f"[yellow]⚠ Already downloaded:[/] {filename}")
|
|
94
|
+
paths.append(target_path)
|
|
95
|
+
continue
|
|
96
|
+
|
|
97
|
+
# Prepare progress bar for this file
|
|
98
|
+
total = int(requests.head(url).headers.get("Content-Length", 0))
|
|
99
|
+
total_bytes += total
|
|
100
|
+
task_id = progress.add_task("Downloading", filename=filename, total=total)
|
|
101
|
+
future = executor.submit(_download_file, url, target_path, progress, task_id)
|
|
102
|
+
futures.append(future)
|
|
103
|
+
|
|
104
|
+
# Wait for all downloads to finish
|
|
105
|
+
for future in as_completed(futures):
|
|
106
|
+
try:
|
|
107
|
+
result = future.result()
|
|
108
|
+
paths.append(result)
|
|
109
|
+
except Exception as e:
|
|
110
|
+
console.print(f"[red]Error downloading file:[/] {e}")
|
|
111
|
+
|
|
112
|
+
duration = time.time() - start_time
|
|
113
|
+
size_mb = total_bytes / (1024 ** 2)
|
|
114
|
+
|
|
115
|
+
# Final summary
|
|
116
|
+
console.print(f"\n[green]📥 {len(paths)} file(s) downloaded[/]")
|
|
117
|
+
console.print(f"[cyan]💾 Total size:[/] {size_mb:.1f} MB")
|
|
118
|
+
console.print(f"[cyan]⏱ Duration:[/] {duration:.1f} seconds")
|
|
119
|
+
|
|
120
|
+
return paths
|
discogs/extractor.py
ADDED
|
@@ -0,0 +1,92 @@
|
|
|
1
|
+
# discogs/extractor.py
|
|
2
|
+
|
|
3
|
+
import gzip
|
|
4
|
+
from pathlib import Path
|
|
5
|
+
from rich.console import Console
|
|
6
|
+
from rich.progress import Progress, SpinnerColumn, BarColumn, TimeElapsedColumn, TextColumn
|
|
7
|
+
|
|
8
|
+
console = Console() # Global console instance for consistent output
|
|
9
|
+
|
|
10
|
+
def extract_gz(gz_path: Path, delete_original: bool = False) -> Path:
|
|
11
|
+
"""
|
|
12
|
+
Extracts a single .gz file into its original XML format.
|
|
13
|
+
Optionally deletes the .gz file after extraction.
|
|
14
|
+
"""
|
|
15
|
+
if gz_path.suffix != ".gz":
|
|
16
|
+
raise ValueError("File is not a .gz file")
|
|
17
|
+
|
|
18
|
+
xml_path = gz_path.with_suffix("") # Remove ".gz" to get .xml filename
|
|
19
|
+
total_size = gz_path.stat().st_size
|
|
20
|
+
|
|
21
|
+
# Display progress bar while extracting
|
|
22
|
+
with Progress(
|
|
23
|
+
SpinnerColumn(),
|
|
24
|
+
TextColumn("[progress.description]{task.description}"),
|
|
25
|
+
BarColumn(),
|
|
26
|
+
"[progress.percentage]{task.percentage:>3.1f}%",
|
|
27
|
+
"•",
|
|
28
|
+
TimeElapsedColumn()
|
|
29
|
+
) as progress:
|
|
30
|
+
task = progress.add_task(f"Extracting {gz_path.name}", total=total_size)
|
|
31
|
+
|
|
32
|
+
with gzip.open(gz_path, 'rb') as f_in, open(xml_path, 'wb') as f_out:
|
|
33
|
+
while True:
|
|
34
|
+
chunk = f_in.read(1024 * 1024) # Read in 1MB chunks
|
|
35
|
+
if not chunk:
|
|
36
|
+
break
|
|
37
|
+
f_out.write(chunk)
|
|
38
|
+
progress.update(task, advance=len(chunk))
|
|
39
|
+
|
|
40
|
+
console.print(f"[green]✔ Extracted:[/] {xml_path}")
|
|
41
|
+
|
|
42
|
+
# Optionally remove the original .gz file after extraction
|
|
43
|
+
if delete_original:
|
|
44
|
+
gz_path.unlink()
|
|
45
|
+
console.print(f"[yellow]🗑 Deleted original:[/] {gz_path}")
|
|
46
|
+
|
|
47
|
+
return xml_path
|
|
48
|
+
|
|
49
|
+
def extract_gz_files(files: list[Path], delete_original: bool = False) -> list[Path]:
|
|
50
|
+
"""
|
|
51
|
+
Extracts multiple .gz files in sequence.
|
|
52
|
+
Returns a list of extracted XML file paths.
|
|
53
|
+
"""
|
|
54
|
+
return [extract_gz(file, delete_original=delete_original) for file in files]
|
|
55
|
+
|
|
56
|
+
def get_extracted_path(gz_path: Path) -> Path:
|
|
57
|
+
"""
|
|
58
|
+
Returns the path of the extracted XML file for a given .gz file.
|
|
59
|
+
"""
|
|
60
|
+
return gz_path.with_suffix("")
|
|
61
|
+
|
|
62
|
+
def extract_interactively():
|
|
63
|
+
"""
|
|
64
|
+
Prompts user to select .gz files for extraction.
|
|
65
|
+
"""
|
|
66
|
+
from rich.prompt import Prompt
|
|
67
|
+
from discogs.config import get_download_dir
|
|
68
|
+
from discogs.utils import open_folder
|
|
69
|
+
|
|
70
|
+
download_dir = get_download_dir()
|
|
71
|
+
dataset_dir = download_dir / "Datasets"
|
|
72
|
+
|
|
73
|
+
gz_files = list(dataset_dir.rglob("*.gz"))
|
|
74
|
+
if not gz_files:
|
|
75
|
+
console.print("[red]No .gz files found to extract.[/red]")
|
|
76
|
+
return
|
|
77
|
+
|
|
78
|
+
console.print("[bold]Select .gz file to extract:[/bold]")
|
|
79
|
+
for i, file in enumerate(gz_files):
|
|
80
|
+
console.print(f"[{i + 1}] {file.relative_to(download_dir)}")
|
|
81
|
+
|
|
82
|
+
choice = Prompt.ask("Enter number", default="1")
|
|
83
|
+
try:
|
|
84
|
+
idx = int(choice.strip()) - 1
|
|
85
|
+
if 0 <= idx < len(gz_files):
|
|
86
|
+
file = gz_files[idx]
|
|
87
|
+
extract_gz(file)
|
|
88
|
+
open_folder(file.parent)
|
|
89
|
+
else:
|
|
90
|
+
console.print("[red]Invalid selection.[/red]")
|
|
91
|
+
except:
|
|
92
|
+
console.print("[red]Invalid input.[/red]")
|