open-standards/aac/aac-rs/tools/pagesets/pslib/data.py

"""Loading: the design sources, the cited word lists, and the app's own data.

Everything is read from the repository, so a build is reproducible offline.
"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from functools import cached_property
from pathlib import Path

try:
    import yaml
except ImportError as exc:  # pragma: no cover
    raise SystemExit("tools/pagesets needs PyYAML: python3 -m pip install pyyaml") from exc

HERE = Path(__file__).resolve().parent.parent  # tools/pagesets
AAC_RS = HERE.parent.parent  # aac/aac-rs
CORE_SRC = AAC_RS / "crates" / "aac-core" / "src"
WEB_SRC = AAC_RS / "crates" / "aac-web" / "src"
OUT = AAC_RS / "data" / "pagesets"

# Part of speech (lexicon) -> the app's word class (aac-core obf.rs names).
WORD_CLASS = {
    "pronoun": "pronoun",
    "determiner": "determiner",
    "verb": "verb",
    "adjective": "adjective",
    "adverb": "adjective",
    "core_noun": "core_noun",
    "fringe_noun": "fringe_noun",
    "preposition": "preposition",
    "question": "question",
    "negation": "negation",
    "social": "social",
    "interjection": "interjection",
    "phrase": "social",
    "conjunction": "conjunction",
}

# Fitzgerald role by part of speech, unless the lexicon says otherwise.
DEFAULT_ROLE = {
    "pronoun": "people",
    "determiner": "little",
    "verb": "action",
    "adjective": "describe",
    "adverb": "describe",
    "core_noun": "thing",
    "fringe_noun": "thing",
    "preposition": "little",
    "conjunction": "little",
    "question": "question",
    "negation": "negation",
    "social": "social",
    "interjection": "social",
    "phrase": "social",
}

# Our Fitzgerald Key palette: (background, border), light tints of the hues
# aac-core gives each word class (WordCategory::fitzgerald, the app's
# defaults.rs): yellow people, green actions, blue describing, orange things,
# pink social, white questions, brown little words (prepositions,
# conjunctions, articles: the key's "other" colour), red negation. Text is
# black on every tint (contrast 13:1 or more).
PALETTE = {
    "people": ("#FFF4C2", "#A67C00"),
    "action": ("#DDF3D8", "#2E7D32"),
    "describe": ("#DCEBFF", "#1565C0"),
    "thing": ("#FFE3C4", "#C25E00"),
    "social": ("#FCE1EC", "#AD1457"),
    "question": ("#FFFFFF", "#6A1B9A"),
    "little": ("#EFE6E1", "#6D4C41"),
    "negation": ("#FFD9D9", "#C62828"),
    "control": ("#E0E0E0", "#37474F"),
}
TEXT_COLOR = "#000000"

CLITICS = {"'s", "'t", "'m", "'re", "'ll", "'ve", "'d"}
STEMS = {
    "don": "don't", "didn": "didn't", "doesn": "doesn't", "isn": "isn't", "wasn": "wasn't",
    "couldn": "couldn't", "wouldn": "wouldn't", "haven": "haven't", "aren": "aren't",
    "shouldn": "shouldn't", "hasn": "hasn't", "weren": "weren't", "ain": "ain't", "won": "won't",
}
MERGE = {"ok": "okay", "gonna": "going", "wanna": "want", "gotta": "got"}
EXCLUDE = {
    # disfluencies
    "uh", "um", "huh", "hmm", "mm", "ah",
    # titles written as abbreviations
    "mr.", "mrs.", "ms.", "dr.",
    # profanity: not proposed by a default page set (a person adds their own)
    "fuck", "fucking", "shit", "damn", "hell", "bitch", "ass", "bastard", "crap", "goddamn",
}


def load_yaml(path: Path):
    with open(path, encoding="utf-8") as f:
        return yaml.safe_load(f)


def read_list(path: Path) -> list[str]:
    return [l.strip() for l in open(path, encoding="utf-8") if l.strip() and not l.startswith("#")]


@dataclass
class Entry:
    word: str
    pos: str
    role: str
    sense: str | None = None
    forms: dict = field(default_factory=dict)
    proxy: str | None = None
    share: float | None = None
    # The key is used in another sense than the app's glosses of it (aac-core
    # Cell::sense): English words only, hidden in place in other languages.
    unglossed: bool = False

    @property
    def word_class(self) -> str | None:
        return WORD_CLASS[self.pos]


class Data:
    """Everything a build reads."""

    def __init__(self):
        self.lexicon = self._lexicon()

    # ── design sources ───────────────────────────────────────────────────
    def _lexicon(self) -> dict[str, Entry]:
        raw = load_yaml(HERE / "data" / "lexicon.en.yaml")["words"]
        out = {}
        for word, spec in raw.items():
            if not isinstance(word, str):
                raise SystemExit(f"lexicon key {word!r} is not text: quote it (YAML reads yes/no/on/off as booleans)")
            pos = spec["pos"]
            if pos not in WORD_CLASS:
                raise SystemExit(f"lexicon: {word}: unknown pos {pos}")
            if " " in word and not spec.get("proxy"):
                raise SystemExit(f"lexicon: {word!r} has more than one word: give it a proxy token")
            out[word] = Entry(
                word, pos, spec.get("role", DEFAULT_ROLE[pos]), spec.get("sense"), spec.get("forms", {}),
                spec.get("proxy"), spec.get("share"), bool(spec.get("unglossed")),
            )
        return out

    def entry(self, word: str) -> Entry:
        if word not in self.lexicon:
            raise SystemExit(f"{word!r} is used by a set but is not in data/lexicon.en.yaml")
        return self.lexicon[word]

    @cached_property
    def growing_core(self) -> dict:
        return load_yaml(HERE / "data" / "growing-core.yaml")

    @cached_property
    def fixed_paths(self) -> dict:
        return load_yaml(HERE / "data" / "fixed-paths.yaml")

    # ── cited sources ────────────────────────────────────────────────────
    @cached_property
    def catalog(self) -> dict:
        return load_yaml(HERE / "sources" / "catalog.yaml")

    @cached_property
    def universal_core(self) -> list[str]:
        return read_list(HERE / "sources" / "universal_core.txt")

    @cached_property
    def banajee(self) -> list[str]:
        return read_list(HERE / "sources" / "banajee_2003.txt")

    @cached_property
    def subtitles(self) -> tuple[dict[str, int], int]:
        """Raw subtitle token counts (lower-case) and the corpus total."""
        counts: dict[str, int] = {}
        total = None
        for line in open(HERE / "sources" / "en_subtitles_2018_top5000.tsv", encoding="utf-8"):
            if line.startswith("#"):
                m = re.search(r"all 50,000 lines of en_50k.txt: (\d+)", line)
                if m:
                    total = int(m.group(1))
                continue
            w, n = line.rstrip("\n").split("\t")
            counts[w] = counts.get(w, 0) + int(n)
        assert total, "the subtitle list states its total"
        return counts, total

    @cached_property
    def spoken_ranking(self) -> list[tuple[str, int]]:
        """Words people say, most frequent first: the subtitle list with clitic
        fragments, disfluencies, abbreviations and profanity left out, and
        contraction stems read as the contraction ("don" -> "don't")."""
        counts, _ = self.subtitles
        merged: dict[str, int] = {}
        for w, n in counts.items():
            if w in CLITICS or w in EXCLUDE:
                continue
            w = STEMS.get(w, w)
            w = MERGE.get(w, w)
            merged[w] = merged.get(w, 0) + n
        return sorted(merged.items(), key=lambda x: (-x[1], x[0]))

    @cached_property
    def spoken_rank(self) -> dict[str, int]:
        return {w: i + 1 for i, (w, _) in enumerate(self.spoken_ranking)}

    def spoken_p(self, token: str) -> float:
        """Relative frequency of a single token in the subtitles. A token
        outside the first 5,000 gets half the 5,000th's count: rare, not
        never said."""
        counts, total = self.subtitles
        t = token.lower()
        t = {v: k for k, v in STEMS.items()}.get(t, t)
        n = counts.get(t, 0) + sum(c for k, c in counts.items() if MERGE.get(k) == t)
        return (n or min(counts.values()) / 2) / total

    def weight(self, entry: "Entry", aac_boost: float, phrase_share: float) -> float:
        """How often a word is likely to be chosen: spoken frequency, times
        (1 + aac_boost) for a word in Universal Core or Banajee et al. 2003."""
        if entry.proxy:
            share = entry.share if entry.share is not None else (phrase_share if entry.pos == "phrase" else 1.0)
            p = self.spoken_p(entry.proxy) * share
        else:
            p = self.spoken_p(entry.word)
        return p * (1 + aac_boost if self.in_aac_lists(entry.word) else 1)

    def in_aac_lists(self, word: str) -> list[str]:
        out = []
        uc = {w.lower() for w in self.universal_core}
        if word.lower() in uc or (word == "all done" and "finished" in uc):
            out.append("universal_core")
        if word.lower() in {w.lower() for w in self.banajee}:
            out.append("banajee_2003")
        return out

    # ── the app's own data ───────────────────────────────────────────────
    @cached_property
    def translations(self) -> dict[str, dict[str, str]]:
        return json.loads((CORE_SRC / "translations_data.json").read_text(encoding="utf-8"))

    def glosses(self, word: str) -> dict[str, str]:
        """translations_for() in aac-core: exact key, then lower-case."""
        return self.translations.get(word) or self.translations.get(word.lower()) or {}

    @cached_property
    def audits(self) -> list[dict]:
        return [json.loads(p.read_text(encoding="utf-8")) for p in sorted(CORE_SRC.glob("gloss_audit_*.json"))]

    @cached_property
    def symbols(self) -> dict:
        return json.loads((CORE_SRC / "symbols" / "index.json").read_text(encoding="utf-8"))

    @cached_property
    def nav_labels(self) -> dict:
        return json.loads((CORE_SRC / "nav_labels.json").read_text(encoding="utf-8"))

    @cached_property
    def speech_lexicon(self) -> dict:
        return json.loads((CORE_SRC / "speech_lexicon.json").read_text(encoding="utf-8"))

    @cached_property
    def app_vocabulary(self) -> dict[str, tuple[str, int, str]]:
        """vocabulary.rs: word (lower-case) -> (app class name, rank, level).
        Like word_class_of(), the later class wins for a word listed twice."""
        src = (CORE_SRC / "vocabulary.rs").read_text(encoding="utf-8")
        pat = re.compile(r'\("([^"]+)",\s*WordCategory::(\w+),\s*(\d+),\s*VocabularyLevel::(\w+),')
        names = {
            "Pronoun": "pronoun", "Verb": "verb", "Adjective": "adjective", "Social": "social",
            "Question": "question", "Preposition": "preposition", "Determiner": "determiner",
            "CoreNoun": "core_noun", "FringeNoun": "fringe_noun", "Interjection": "interjection",
            "Negation": "negation", "Conjunction": "conjunction",
        }
        order = list(names.values())
        out: dict[str, tuple[str, int, str]] = {}
        # word_class_of() inserts class by class in WORD_CLASSES order, so a
        # word in two classes ends with the later class.
        rows = [(w, names[c], int(r), lvl) for w, c, r, lvl in pat.findall(src)]
        for cls in order:
            for w, c, r, lvl in rows:
                if c == cls:
                    out[w.lower()] = (c, r, lvl)
        return out

    @cached_property
    def languages(self) -> list[dict]:
        """The app's 45 languages (i18n.rs), with their packs (data/lang)."""
        src = (CORE_SRC / "i18n.rs").read_text(encoding="utf-8")
        pat = re.compile(r'lang\("(\w+)",\s*"([\w-]+)",\s*"([^"]*)",\s*"([^"]*)"')
        langs = [{"code": c, "locale": l, "native": n, "english": e} for c, l, n, e in pat.findall(src)]
        packs = json.loads((AAC_RS / "data" / "lang" / "index.json").read_text(encoding="utf-8"))
        for lang in langs:
            pack = packs.get("packs", {}).get(lang["code"])
            lang["pack_quality"] = pack.get("quality") if pack else None
            lang["core_status"] = "ValidatedNativeCorpus" if lang["code"] == "en" else "GlossOnly"
        return langs