Files

193 lines
6.8 KiB
Python
Raw Permalink Normal View History

import os
import re
import sys
import time
import urllib.parse
from bs4 import BeautifulSoup
from curl_cffi import requests
import pandas as pd
from rapidfuzz import fuzz
from concurrent.futures import ThreadPoolExecutor, as_completed
import threading
DATA_DIR = "data"
CHECKPOINT_DIR = os.path.join(DATA_DIR, "checkpoints")
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
MASTER_PATH = os.path.join(DATA_DIR, "master_artists_17k.csv")
CSPC_PATH = os.path.join(DATA_DIR, "chartmasters_audited.csv")
CHECKPOINT_FILE = os.path.join(CHECKPOINT_DIR, "sales_pipeline_progress.csv")
FINAL_OUTPUT = os.path.join(DATA_DIR, "master_sales_17k_verified.csv")
checkpoint_lock = threading.Lock()
def normalize_text(text: str) -> str:
if not text:
return ""
cleaned = re.sub(r"\(.*?\)|\[.*?\]", "", str(text))
cleaned = re.sub(r"[^\w\s]", "", cleaned)
return cleaned.strip().lower()
def extract_units_from_title(title_text: str) -> int:
text = (title_text or "").upper().strip()
mult_match = re.search(r"(\d+)\s*X", text)
multiplier = int(mult_match.group(1)) if mult_match else 1
if "DIAMOND" in text:
return max(multiplier, 1) * 10_000_000
elif "PLATINUM" in text or "MULTI-PLATINUM" in text:
return multiplier * 1_000_000
elif "GOLD" in text:
return 500_000
return 0
def is_strict_match(target_name: str, candidate_artist: str, candidate_title: str) -> bool:
c_target = normalize_text(target_name)
if not c_target or len(c_target) < 2:
return False
c_cand_artist = normalize_text(candidate_artist)
c_cand_title = normalize_text(candidate_title)
if c_target == c_cand_artist:
return True
if re.search(rf"\b{re.escape(c_target)}\b", c_cand_artist):
return True
if re.search(rf"\b(feat|ft|featuring|with)\s+.*?\b{re.escape(c_target)}\b", candidate_title.lower()):
return True
if len(c_target) > 5 and fuzz.token_sort_ratio(c_target, c_cand_artist) >= 88:
return True
return False
def fetch_riaa(artist_name: str) -> dict:
target_clean = normalize_text(artist_name)
if not target_clean or len(target_clean) < 2:
return {"units": 0, "awards": 0, "provenance": "UNVERIFIED_NO_RECORDS"}
encoded = urllib.parse.quote_plus(target_clean)
url = f"https://www.riaa.com/gold-platinum/?tab_active=default-award&se={encoded}#search_section"
try:
resp = requests.get(url, impersonate="chrome120", timeout=12)
if resp.status_code != 200:
return {"units": 0, "awards": 0, "provenance": "UNVERIFIED_NO_RECORDS"}
soup = BeautifulSoup(resp.text, "html.parser")
table = soup.find("table")
if not table:
return {"units": 0, "awards": 0, "provenance": "UNVERIFIED_NO_RECORDS"}
rows = table.find_all("tr")
total_units = 0
award_count = 0
for r in rows[1:]:
cells = r.find_all("td")
if len(cells) < 3:
continue
row_artist = cells[1].get_text(strip=True)
row_title = cells[2].get_text(strip=True)
if not is_strict_match(artist_name, row_artist, row_title):
continue
img = cells[0].find("img")
units = extract_units_from_title(img.get("title", "") if img else "")
if units > 0:
total_units += units
award_count += 1
if award_count > 0:
return {"units": total_units, "awards": award_count, "provenance": "VERIFIED_RIAA"}
except Exception:
pass
return {"units": 0, "awards": 0, "provenance": "UNVERIFIED_NO_RECORDS"}
def process_artist(row, cspc_dict):
a_id = row["artist_id"]
a_name = str(row["artist_name"])
c_name = normalize_text(a_name)
# Tier 1: CSPC
if c_name in cspc_dict:
c_data = cspc_dict[c_name]
return {
"artist_id": a_id,
"artist_name": a_name,
"audited_sales_units": float(c_data["total_eas_cspc"]),
"awards_count": None,
"provenance": "VERIFIED_CSPC"
}
# Tier 2: RIAA
riaa_res = fetch_riaa(a_name)
units = riaa_res["units"]
return {
"artist_id": a_id,
"artist_name": a_name,
"audited_sales_units": units if units > 0 else None,
"awards_count": riaa_res["awards"] if riaa_res["awards"] > 0 else None,
"provenance": riaa_res["provenance"]
}
def run_pipeline(max_workers=5):
print(f"Initializing Multi-Threaded Sales Pipeline ({max_workers} Workers)...")
master_df = pd.read_csv(MASTER_PATH)
cspc_df = pd.read_csv(CSPC_PATH)
cspc_dict = {normalize_text(r["artist_name"]): r for _, r in cspc_df.iterrows()}
if os.path.exists(CHECKPOINT_FILE):
progress_df = pd.read_csv(CHECKPOINT_FILE)
processed_ids = set(progress_df["artist_id"].dropna())
records = progress_df.to_dict("records")
print(f"Resuming from checkpoint: {len(processed_ids):,} artists already processed.")
else:
records = []
processed_ids = set()
queue_df = master_df[~master_df["artist_id"].isin(processed_ids)]
total_remaining = len(queue_df)
total_catalog = len(master_df)
print(f"Running on remaining {total_remaining:,} artists out of {total_catalog:,}...\n")
processed_count = len(processed_ids)
with ThreadPoolExecutor(max_workers=max_workers) as executor:
future_to_artist = {
executor.submit(process_artist, row, cspc_dict): row
for _, row in queue_df.iterrows()
}
for future in as_completed(future_to_artist):
res = future.result()
records.append(res)
processed_count += 1
if res["audited_sales_units"] is not None and res["audited_sales_units"] > 0:
print(f"[{processed_count:,}/{total_catalog:,}] {res['artist_name']:<22} -> {res['audited_sales_units']:>12,f} ({res['provenance']})")
# Save checkpoint every 50 artists
if processed_count % 50 == 0:
with checkpoint_lock:
pd.DataFrame(records).to_csv(CHECKPOINT_FILE, index=False)
# Final Rank & Export
df_out = pd.DataFrame(records)
has_sales = df_out["audited_sales_units"].notna()
df_out["sales_rank"] = None
df_out.loc[has_sales, "sales_rank"] = (
df_out.loc[has_sales, "audited_sales_units"]
.rank(ascending=False, method="min")
.astype(int)
)
df_out.to_csv(FINAL_OUTPUT, index=False)
print(f"\nPipeline successfully finished! Final catalog saved to: {FINAL_OUTPUT}")
if __name__ == "__main__":
run_pipeline(max_workers=5)