Files
dwroller/scripts/fix-talent-titles.py

127 lines
4.6 KiB
Python

#!/usr/bin/env python3
"""
Fix garbled OCR talent titles in the rules DB.
- Normalise mixed-case OCR titles to proper Title Case
- Set source/source_abbr/category correctly for all talent entries
- Remove duplicate sanitized entries where a clean 'talent_*' version exists
"""
import os, re, mysql.connector
from pathlib import Path
REPO = Path(__file__).resolve().parent.parent
def load_env_file(path):
if not path.exists():
return
for line in path.read_text().splitlines():
line = line.strip()
if not line or line.startswith('#') or '=' not in line:
continue
key, value = line.split('=', 1)
os.environ.setdefault(key.strip(), value.strip().strip('"').strip("'"))
load_env_file(REPO / 'database' / '.env')
load_env_file(REPO / '.env')
if not os.environ.get('DB_PASSWORD'):
raise RuntimeError('DB_PASSWORD must be set in database/.env or .env')
DB = dict(
host=os.environ.get('DB_HOST', '192.168.1.113'),
port=int(os.environ.get('DB_PORT', '3307')),
user=os.environ.get('DB_USER', 'deathwatch'),
password=os.environ['DB_PASSWORD'],
database=os.environ.get('DB_NAME', 'deathwatch'),
)
def normalise_title(raw):
"""Convert OCR-garbled title to proper Title Case."""
# Remove excess whitespace
s = re.sub(r'\s+', ' ', raw).strip()
# Lowercase everything first, then title-case
words = s.lower().split()
# Keep small words lowercase unless first word
small = {'a','an','the','of','to','in','for','and','or','but','with','at','by','from'}
result = []
for i, w in enumerate(words):
if i == 0 or w not in small:
# Handle hyphenated words
result.append('-'.join(p.capitalize() for p in w.split('-')))
else:
result.append(w)
return ' '.join(result)
db = mysql.connector.connect(**DB)
c = db.cursor()
# 1. Get all sanitized talent entries (have "Prerequisites:" in content)
c.execute("""SELECT id, rule_id, title FROM rules
WHERE source = 'sanitized' AND content LIKE 'Prerequisites:%'""")
sanitized_talents = c.fetchall()
print(f"Found {len(sanitized_talents)} sanitized talent entries")
# 2. Get existing clean talent rule_ids (from our earlier import)
c.execute("SELECT rule_id, title FROM rules WHERE source = 'Core Rulebook' AND category = 'talents'")
clean_talents = {row[0]: row[1] for row in c.fetchall()}
print(f"Found {len(clean_talents)} clean talent entries")
fixed = 0
deleted = 0
for (db_id, rule_id, title) in sanitized_talents:
clean_title = normalise_title(title)
expected_rid = 'talent_' + re.sub(r'[^a-z0-9]', '_', clean_title.lower())
# Remove trailing underscores
expected_rid = re.sub(r'_+$', '', expected_rid)
expected_rid = re.sub(r'_+', '_', expected_rid)
if expected_rid in clean_talents:
# We already have a clean version - delete the sanitized duplicate
c.execute("DELETE FROM rules WHERE id = %s", (db_id,))
deleted += 1
else:
# No clean version exists - fix this entry's title/source/category
c.execute("""UPDATE rules
SET title = %s, source = 'Core Rulebook', source_abbr = 'CR', category = 'talents'
WHERE id = %s""",
(clean_title, db_id))
# Also add a proper rule_id if missing
new_rid = expected_rid
# Check if rule_id already conflicts
c.execute("SELECT COUNT(*) FROM rules WHERE rule_id = %s AND id != %s", (new_rid, db_id))
if c.fetchone()[0] == 0:
c.execute("UPDATE rules SET rule_id = %s WHERE id = %s", (new_rid, db_id))
fixed += 1
db.commit()
print(f"Fixed: {fixed} talent titles, Deleted: {deleted} duplicates")
# 3. Also fix remaining garbled-title talents (non-sanitized, left-over OCR artifacts)
c.execute("""SELECT id, title FROM rules
WHERE category = 'talents' AND title REGEXP '[A-Z]{2,}[a-z]' """)
garbled = c.fetchall()
print(f"\nFixing {len(garbled)} remaining garbled-title talent entries...")
for (db_id, title) in garbled:
clean = normalise_title(title)
if clean != title:
c.execute("UPDATE rules SET title = %s WHERE id = %s", (clean, db_id))
print(f" '{title}''{clean}'")
db.commit()
# 4. Final stats
c.execute("SELECT source, COUNT(*) FROM rules GROUP BY source ORDER BY COUNT(*) DESC")
print("\n=== Final rule counts by source ===")
total = 0
for src, cnt in c.fetchall():
print(f" {src}: {cnt}")
total += cnt
print(f" TOTAL: {total}")
c.execute("SELECT category, COUNT(*) FROM rules GROUP BY category ORDER BY COUNT(*) DESC")
print("\n=== By category ===")
for cat, cnt in c.fetchall():
print(f" {cat}: {cnt}")
c.close()
db.close()