UnheardCity/ackit/loctext.py

"""Menu and UI text: the game's localization files BmGame/Localization/<LANG>/<File>.<lang>
(UE3 .ini format; docs/formats/text.md).

Files are UTF-16LE with a BOM and CRLF line ends (two small ones are plain ASCII). Lines are `[Section]`,
`Key=Value` (the value may be in double quotes) or comments (`;`). A key can repeat inside a section, so an
entry is identified by (file, section, key, occurrence). Writing changes only the value part of the edited
lines; every other byte of the file (BOM, encoding, quotes, line ends, comments) stays as it was.
"""
import os
import re
from collections import namedtuple

# key: "<LANG>:<File>:<Section>:<Key>" + ("#n" for the n-th repeat, n >= 2)
Entry = namedtuple("Entry", "key lang file section name occurrence value quoted line")


class TextError(ValueError):
    pass


def _decode(raw):
    if raw[:2] == b"\xff\xfe":
        return raw[2:].decode("utf-16-le"), "utf-16"
    if raw[:3] == b"\xef\xbb\xbf":
        return raw[3:].decode("utf-8"), "utf-8"
    return raw.decode("cp1252"), "cp1252"


def _encode(text, enc):
    if enc == "utf-16":
        return b"\xff\xfe" + text.encode("utf-16-le")
    if enc == "utf-8":
        return b"\xef\xbb\xbf" + text.encode("utf-8")
    return text.encode("cp1252")


_KV = re.compile(r"^([^=;\[\s][^=]*?)=(.*)$")


def parse(raw, lang, fname):
    """Entries of one file's bytes."""
    text, _ = _decode(raw)
    out = []
    section = ""
    seen = {}
    for i, line in enumerate(re.split(r"(?<=\n)", text)):
        body = line.rstrip("\r\n")
        s = body.strip()
        if s.startswith("[") and s.endswith("]"):
            section = s[1:-1]
            continue
        m = _KV.match(body)
        if not m or s.startswith(";"):
            continue
        name, value = m.group(1), m.group(2)
        quoted = len(value) >= 2 and value[0] == '"' and value[-1] == '"'
        if quoted:
            value = value[1:-1]
        k = (section.lower(), name.lower())
        seen[k] = seen.get(k, 0) + 1
        n = seen[k]
        key = "%s:%s:%s:%s" % (lang, fname, section, name) + ("#%d" % n if n > 1 else "")
        out.append(Entry(key, lang, fname, section, name, n, value, quoted, i))
    return out


def encoding(raw):
    """Encoding of a file's bytes: 'utf-16', 'utf-8' or 'cp1252'."""
    return "utf-16" if raw[:2] == b"\xff\xfe" else "utf-8" if raw[:3] == b"\xef\xbb\xbf" else "cp1252"


def check_value(value, enc, fname):
    """Raise TextError if `value` cannot be written into a file of encoding `enc` (what `write` refuses)."""
    if "\r" in value or "\n" in value:
        raise TextError("a menu text cannot contain line breaks (use the game's own markers, e.g. \\n)")
    try:
        _encode(value, enc)
    except UnicodeEncodeError:
        raise TextError("%s is stored as %s and cannot hold these characters" % (fname, enc)) from None


def write(raw, lang, fname, values):
    """New file bytes with {entry key: new value} applied (keys not in the file raise TextError)."""
    text, enc = _decode(raw)
    lines = re.split(r"(?<=\n)", text)
    by_key = {e.key: e for e in parse(raw, lang, fname)}
    for key, value in values.items():
        e = by_key.get(key)
        if e is None:
            raise TextError("no text %s in %s" % (key, fname))
        check_value(value, enc, fname)
        line = lines[e.line]
        end = line[len(line.rstrip("\r\n")):]
        head = line[:line.index("=") + 1]
        lines[e.line] = head + ('"%s"' % value if e.quoted else value) + end
    out = "".join(lines)
    try:
        return _encode(out, enc)
    except UnicodeEncodeError:
        raise TextError("%s is stored as %s and cannot hold these characters" % (fname, enc)) from None


def files(loc_root, lang):
    """{file name: path} of one language folder (case-insensitive): localization files, without the
    Version*.int build IDs (not UI text)."""
    d = next((os.path.join(loc_root, x) for x in os.listdir(loc_root) if x.upper() == lang.upper()), None)
    if not d:
        return {}
    return {f: os.path.join(d, f) for f in sorted(os.listdir(d))
            if os.path.splitext(f)[1].lower() == "." + lang.lower() and not f.lower().startswith("version")}


def languages(loc_root):
    return sorted(x.upper() for x in os.listdir(loc_root) if os.path.isdir(os.path.join(loc_root, x)))


def load(loc_root, lang):
    """All entries of a language: [Entry]."""
    out = []
    for f, p in files(loc_root, lang).items():
        out += parse(open(p, "rb").read(), lang.upper(), f)
    return out