cogcore/data/document.py

"""Document dataclass for corpus manifests."""
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Mapping


@dataclass(frozen=True)
class CorpusDocument:
    id: str
    text: str
    lang: str = "unknown"
    domain: str = "general"
    source_type: str = "local"
    quality: float = 1.0
    split: str | None = None
    tags: tuple[str, ...] = field(default_factory=tuple)
    meta: Mapping[str, Any] = field(default_factory=dict)

    @classmethod
    def from_mapping(cls, obj: Mapping[str, Any]) -> "CorpusDocument":
        return cls(
            id=str(obj.get("id", "")),
            text=str(obj.get("text", "")),
            lang=str(obj.get("lang", "unknown")),
            domain=str(obj.get("domain", "general")),
            source_type=str(obj.get("source_type", "local")),
            quality=float(obj.get("quality", 1.0)),
            split=obj.get("split"),
            tags=tuple(obj.get("tags", ()) or ()),
            meta=dict(obj.get("meta", {}) or {}),
        )

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "text": self.text,
            "lang": self.lang,
            "domain": self.domain,
            "source_type": self.source_type,
            "quality": self.quality,
            "split": self.split,
            "tags": list(self.tags),
            "meta": dict(self.meta),
        }

    def with_text(self, text: str) -> "CorpusDocument":
        data = self.to_dict()
        data["text"] = text
        return CorpusDocument.from_mapping(data)

    def with_split(self, split: str) -> "CorpusDocument":
        data = self.to_dict()
        data["split"] = split
        return CorpusDocument.from_mapping(data)