src/wvy_experimental/data.py

from __future__ import annotations

import csv
import json
from collections.abc import Iterator
from pathlib import Path

SUPPORTED_EXTENSIONS = {".txt", ".md", ".jsonl", ".json", ".csv", ".parquet"}


def discover_files(path: str | Path) -> list[Path]:
    root = Path(path).expanduser().resolve()
    if root.is_file():
        files = [root]
    elif root.is_dir():
        files = [item for item in root.rglob("*") if item.is_file()]
    else:
        raise FileNotFoundError(root)
    supported = sorted(item for item in files if item.suffix.lower() in SUPPORTED_EXTENSIONS)
    if not supported:
        raise ValueError(f"No supported data files found under {root}")
    return supported


def iter_documents(files: list[Path], text_field: str = "text") -> Iterator[str]:
    for path in files:
        suffix = path.suffix.lower()
        if suffix in {".txt", ".md"}:
            text = path.read_text(encoding="utf-8", errors="replace")
            if text.strip():
                yield text
        elif suffix == ".jsonl":
            with path.open("r", encoding="utf-8", errors="replace") as handle:
                for line_number, line in enumerate(handle, 1):
                    if not line.strip():
                        continue
                    value = json.loads(line)
                    text = value.get(text_field) if isinstance(value, dict) else None
                    if isinstance(text, str) and text.strip():
                        yield text
                    elif text is None:
                        raise ValueError(f"Missing string field {text_field!r} in {path}:{line_number}")
        elif suffix == ".json":
            value = json.loads(path.read_text(encoding="utf-8", errors="replace"))
            records = value if isinstance(value, list) else [value]
            for index, record in enumerate(records):
                text = record.get(text_field) if isinstance(record, dict) else None
                if isinstance(text, str) and text.strip():
                    yield text
                elif text is None:
                    raise ValueError(f"Missing string field {text_field!r} in {path}, record {index}")
        elif suffix == ".csv":
            with path.open("r", encoding="utf-8", errors="replace", newline="") as handle:
                for index, record in enumerate(csv.DictReader(handle), 2):
                    text = record.get(text_field)
                    if isinstance(text, str) and text.strip():
                        yield text
                    elif text is None:
                        raise ValueError(f"Missing field {text_field!r} in {path}:{index}")
        elif suffix == ".parquet":
            try:
                import pyarrow.parquet as pq
            except ImportError as error:
                raise RuntimeError("Parquet scanning requires pyarrow") from error
            table = pq.read_table(path, columns=[text_field])
            for text in table[text_field].to_pylist():
                if isinstance(text, str) and text.strip():
                    yield text