#!/usr/bin/env python3 """ Import fruits from 03-data/obstsorten.xml into the OSDB database. Usage: DATABASE_URL=postgres://... python3 scripts/import_fruits.py Idempotent: deletes synonyms and images for each fruit before re-inserting. Upserts the fruit row by osdb_number. """ import glob import os import sys import xml.etree.ElementTree as ET from datetime import date, datetime import psycopg2 DATA_DIR = os.path.join(os.path.dirname(__file__), "..", "03-data") XML_PATH = os.path.join(DATA_DIR, "obstsorten.xml") # typ values that are aggregate categories — coerced with a logged warning AGGREGATE_TYPS = {"x", "bq", "aprpfi", "mirren", "pflzwe", "bo", "k"} TYP_MAP = { "a": "Apfelsorten", "x": "Apfelsorten", "b": "Birnensorten", "bq": "Birnensorten", "q": "Quittensorten", "apr": "Aprikosen", "aprpfi": "Aprikosen", "pfi": "Pfirsiche", "p": "Pfirsiche", "mir": "Mirabellen", "mirren": "Mirabellen", "ren": "Renekloden", "pfl": "Pflaumen", "pflzwe": "Pflaumen", "zwe": "Zwetschen", "k": "Sauerkirschen", "bo": "Brombeeren", "bob": "Brombeeren", "boe": "Erdbeeren", "boh": "Himbeeren", "boj": "Johannisbeeren", "bos": "Stachelbeeren", "wei": "Wein", "w": "Wein", } IMAGE_DIRS = [ ("einzelfruechte", "fruit"), ("blueten", "flower"), ("baume", "tree"), ] def map_typ(typ: str): """Map XML typ code to fruit_type enum value. Returns None for unknowns.""" return TYP_MAP.get(typ) def parse_synonyms(text) -> list: """Split comma-separated synonym text into a clean list.""" if not text: return [] return [s.strip() for s in text.split(",") if s.strip()] def parse_date(text) -> date | None: """Parse YYYYMMDD string to date. Returns None on failure or empty input.""" if not text: return None try: return datetime.strptime(text.strip(), "%Y%m%d").date() except ValueError: return None def find_images(fruit_id: str, data_dir: str) -> list: """ Find all _s0 images for a fruit across the three image directories. Returns list of dicts with keys: filename, data, image_type. """ results = [] for subdir, image_type in IMAGE_DIRS: pattern = os.path.join(data_dir, subdir, f"{fruit_id}_*_s0.*") for path in sorted(glob.glob(pattern)): with open(path, "rb") as f: data = f.read() results.append( { "filename": os.path.basename(path), "data": data, "image_type": image_type, } ) return results def import_fruits(conn, data_dir: str = None) -> dict: """ Import all fruits from obstsorten.xml. Returns counts dict: fruits, synonyms, images, warnings, skipped. """ if data_dir is None: data_dir = DATA_DIR xml_path = os.path.join(data_dir, "obstsorten.xml") tree = ET.parse(xml_path) root = tree.getroot() counts = {"fruits": 0, "synonyms": 0, "images": {"fruit": 0, "flower": 0, "tree": 0}, "warnings": 0, "skipped": 0} warnings = [] with conn.cursor() as cur: for sorte in root.findall("sorte"): osdb_number = (sorte.findtext("id") or "").strip() name = (sorte.findtext("name") or "").strip() typ = (sorte.findtext("typ") or "").strip() synonym_text = sorte.findtext("synonym") added_text = (sorte.findtext("added") or "").strip() if not osdb_number or not name: counts["skipped"] += 1 warnings.append(f"SKIP: missing id or name (id={osdb_number!r})") continue fruit_type = map_typ(typ) if fruit_type is None: counts["skipped"] += 1 warnings.append(f"SKIP: unknown typ={typ!r} for id={osdb_number!r}") counts["warnings"] += 1 continue if typ in AGGREGATE_TYPS: warnings.append(f"COERCE: typ={typ!r} → {fruit_type!r} for id={osdb_number!r}") counts["warnings"] += 1 created_at = parse_date(added_text) synonyms = parse_synonyms(synonym_text) # Upsert fruit — get id back for child inserts if created_at: cur.execute( """ INSERT INTO fruits (name, osdb_number, fruit_type, created_at, updated_at) VALUES (%s, %s, %s, %s, NOW()) ON CONFLICT (osdb_number) DO UPDATE SET name = EXCLUDED.name, fruit_type = EXCLUDED.fruit_type, comment = NULL, updated_at = NOW() RETURNING id """, (name, osdb_number, fruit_type, created_at), ) else: cur.execute( """ INSERT INTO fruits (name, osdb_number, fruit_type, updated_at) VALUES (%s, %s, %s, NOW()) ON CONFLICT (osdb_number) DO UPDATE SET name = EXCLUDED.name, fruit_type = EXCLUDED.fruit_type, comment = NULL, updated_at = NOW() RETURNING id """, (name, osdb_number, fruit_type), ) fruit_id = cur.fetchone()[0] counts["fruits"] += 1 # Idempotent: delete existing synonyms and images cur.execute("DELETE FROM fruit_synonyms WHERE fruit_id = %s", (fruit_id,)) cur.execute("DELETE FROM fruit_images WHERE fruit_id = %s", (fruit_id,)) # Insert synonyms for syn in synonyms: cur.execute( "INSERT INTO fruit_synonyms (fruit_id, synonym) VALUES (%s, %s)", (fruit_id, syn), ) counts["synonyms"] += 1 # Insert images for img in find_images(osdb_number, data_dir): cur.execute( """ INSERT INTO fruit_images (fruit_id, filename, data, image_type, created_at) VALUES (%s, %s, %s, %s, NOW()) """, (fruit_id, img["filename"], psycopg2.Binary(img["data"]), img["image_type"]), ) counts["images"][img["image_type"]] += 1 conn.commit() for w in warnings[:20]: print(f" WARNING: {w}", file=sys.stderr) if len(warnings) > 20: print(f" ... and {len(warnings) - 20} more warnings", file=sys.stderr) return counts def main(): database_url = os.environ.get("DATABASE_URL") if not database_url: print("Error: DATABASE_URL environment variable not set", file=sys.stderr) sys.exit(1) print("Connecting to database...") conn = psycopg2.connect(database_url) try: print("Importing fruits...") counts = import_fruits(conn) finally: conn.close() total_images = sum(counts["images"].values()) print(f"\nDone.") print(f" Fruits upserted: {counts['fruits']}") print(f" Synonyms inserted: {counts['synonyms']}") print(f" Images inserted: {total_images}") print(f" fruit: {counts['images']['fruit']}") print(f" flower: {counts['images']['flower']}") print(f" tree: {counts['images']['tree']}") print(f" Warnings: {counts['warnings']}") print(f" Skipped: {counts['skipped']}") if __name__ == "__main__": main()