webmail-client/deploy/ngrams.py

"""Download and unpack LanguageTool n-gram data (run by `webmail.sh ngrams`).

Runs inside the backend image (any Python 3 will do) with the n-gram folder
mounted at /ngrams; usage: ngrams.py [lang ...]  (default: en de fr es).

Per language: skip it if already installed, check free space, download with
resume (an 8 GB transfer that drops halfway picks up where it stopped), verify
the size, unpack into a temporary folder and move it into place in one step —
LanguageTool never sees half a language — then delete the zip. Python's zipfile
rather than unzip: BusyBox unzip cannot be trusted with an 8 GB ZIP64 archive.
"""

from __future__ import annotations

import os
import shutil
import sys
import time
import urllib.request
import zipfile
from pathlib import Path

BASE = "https://languagetool.org/download/ngram-data/"
FILES = {
    "en": "ngrams-en-20150817.zip",
    "de": "ngrams-de-20150819.zip",
    "fr": "ngrams-fr-20150913.zip",
    "es": "ngrams-es-20150915.zip",
    "nl": "ngrams-nl-20181229.zip",
}
ROOT = Path(os.environ.get("NGRAMS_DIR", "/ngrams"))
CHUNK = 1 << 20
# languagetool.org answers Python's default User-Agent with 403 Forbidden.
HEADERS = {"User-Agent": "Mozilla/5.0 (compatible; webmail-ngrams/1.0)"}


def say(msg: str) -> None:
    print(msg, flush=True)


def gib(n: float) -> str:
    return f"{n / 1024**3:.2f} GB"


def installed(lang: str) -> bool:
    d = ROOT / lang
    return all((d / f"{n}grams").is_dir() for n in (1, 2, 3))


def remote_size(url: str) -> int:
    req = urllib.request.Request(url, method="HEAD", headers=HEADERS)
    with urllib.request.urlopen(req, timeout=60) as resp:
        return int(resp.headers["Content-Length"])


def download(url: str, dest: Path, total: int) -> None:
    have = dest.stat().st_size if dest.exists() else 0
    if have == total:
        say("  already downloaded")
        return
    req = urllib.request.Request(url, headers=HEADERS)
    if have:
        req.add_header("Range", f"bytes={have}-")
        say(f"  resuming at {gib(have)}")
    with (
        urllib.request.urlopen(req, timeout=60) as resp,
        open(dest, "ab" if have else "wb") as out,
    ):
        if have and resp.status != 206:  # the server ignored the range: start over
            out.truncate(0)
            have = 0
        done, last, started = have, -1, time.monotonic()
        while chunk := resp.read(CHUNK):
            out.write(chunk)
            done += len(chunk)
            pct = done * 100 // total
            if pct // 5 != last // 5:
                rate = (done - have) / max(time.monotonic() - started, 0.001)
                say(
                    f"  {pct:3d}%  {gib(done)} of {gib(total)}  ({rate / 1024**2:.1f} MB/s)"
                )
                last = pct
    if dest.stat().st_size != total:
        raise SystemExit(
            f"  download incomplete ({gib(dest.stat().st_size)} of {gib(total)}) — run again to resume"
        )


def unpack(zip_path: Path, lang: str) -> None:
    tmp = ROOT / f".unpacking-{lang}"
    shutil.rmtree(tmp, ignore_errors=True)
    say("  unpacking…")
    with zipfile.ZipFile(zip_path) as zf:
        zf.extractall(tmp)
    # The archive holds "<lang>/1grams …"; tolerate a different top folder.
    src = next((p.parent for p in tmp.rglob("1grams") if p.is_dir()), None)
    if src is None:
        shutil.rmtree(tmp, ignore_errors=True)
        raise SystemExit(
            f"  {zip_path.name} has no 1grams folder — not LanguageTool n-gram data?"
        )
    for p in [src, *src.rglob("*")]:  # readable by the languagetool user (uid 100)
        p.chmod(0o755 if p.is_dir() else 0o644)
    final = ROOT / lang
    shutil.rmtree(final, ignore_errors=True)
    src.rename(final)
    shutil.rmtree(tmp, ignore_errors=True)


def size_of(path: Path) -> int:
    return sum(p.stat().st_size for p in path.rglob("*") if p.is_file())


def main(langs: list[str]) -> None:
    ROOT.mkdir(parents=True, exist_ok=True)
    unknown = [lang for lang in langs if lang not in FILES]
    if unknown:
        raise SystemExit(
            f"no n-gram data for: {' '.join(unknown)} (available: {' '.join(FILES)})"
        )
    for lang in langs:
        say(f"== {lang}")
        if installed(lang):
            say(f"  already installed ({gib(size_of(ROOT / lang))})")
            continue
        url = BASE + FILES[lang]
        total = remote_size(url)
        zip_path = ROOT / FILES[lang]
        have = zip_path.stat().st_size if zip_path.exists() else 0
        # Room for the rest of the zip plus the unpacked copy, which is about
        # twice the zip's size (German: 1.58 GB download, 3.06 GB unpacked).
        need = (total - have) + total * 2
        free = shutil.disk_usage(ROOT).free
        if free < need:
            raise SystemExit(f"  not enough space: need ~{gib(need)}, {gib(free)} free")
        say(f"  downloading {FILES[lang]} ({gib(total)})")
        download(url, zip_path, total)
        unpack(zip_path, lang)
        zip_path.unlink()
        say(f"  installed: {gib(size_of(ROOT / lang))}")
    say(
        "done: "
        + ", ".join(
            f"{lang} ({gib(size_of(ROOT / lang))})" for lang in FILES if installed(lang)
        )
    )


if __name__ == "__main__":
    main(sys.argv[1:] or ["en", "de", "fr", "es"])