253 lines
8.3 KiB
Python
253 lines
8.3 KiB
Python
#!/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(os.path.abspath(__file__)), "..", "03-data")
|
|
|
|
# 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"),
|
|
("osdb/fru", "fruit"),
|
|
("osdb/frucht", "fruit"),
|
|
("blueten", "flower"),
|
|
("osdb/blu", "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"{glob.escape(fruit_id)}_*_s0.*")
|
|
for path in sorted(glob.glob(pattern)):
|
|
basename = os.path.basename(path)
|
|
# Verify the fruit_id is the exact prefix, not just a string prefix of
|
|
# another fruit's id (e.g. "mitschurins" must not steal
|
|
# "mitschurins_fruchtbare_nda_s0.jpg").
|
|
# Convention: basename = {id}_{type_code}_s0.{ext}
|
|
# So stripping the last _-segment before _s0 gives the exact id.
|
|
stem = basename.split("_s0.")[0]
|
|
actual_id = "_".join(stem.split("_")[:-1])
|
|
if actual_id != fruit_id:
|
|
continue
|
|
with open(path, "rb") as f:
|
|
data = f.read()
|
|
results.append(
|
|
{
|
|
"filename": basename,
|
|
"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": {img_type: 0 for _, img_type in IMAGE_DIRS}, "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
|
|
counts["warnings"] += 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 = fruits.comment,
|
|
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 = fruits.comment,
|
|
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)
|
|
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}")
|
|
for img_type, n in counts["images"].items():
|
|
print(f" {img_type}: {n}")
|
|
print(f" Warnings: {counts['warnings']}")
|
|
print(f" Skipped: {counts['skipped']}")
|
|
except Exception as e:
|
|
print(f"Error: {e}", file=sys.stderr)
|
|
conn.rollback()
|
|
sys.exit(1)
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|