xkeyboard-config-test.py.in

#!/usr/bin/env python3

from __future__ import annotations

import argparse
import gzip
import itertools
import multiprocessing
import os
import subprocess
import sys
import xml.etree.ElementTree as ET
from abc import ABCMeta, abstractmethod
from dataclasses import dataclass
from functools import partial
from pathlib import Path
from typing import (
    TYPE_CHECKING,
    ClassVar,
    Iterable,
    Iterator,
    NoReturn,
    Protocol,
    Sequence,
    TextIO,
    TypeVar,
    cast,
)

# TODO: import unconditionnaly Self from typing once we raise Python requirement to 3.11+
if TYPE_CHECKING:
    from typing_extensions import Self

WILDCARD = "*"

DEFAULT_RULES_XML = "@XKB_CONFIG_ROOT@/rules/evdev.xml"

# Meson needs to fill this in so we can call the tool in the buildir.
EXTRA_PATH = "@MESON_BUILD_ROOT@"
os.environ["PATH"] = ":".join(filter(bool, (EXTRA_PATH, os.getenv("PATH"))))

# Environment variable to get the right level of log
os.environ["XKB_LOG_LEVEL"] = "warning"
os.environ["XKB_LOG_VERBOSITY"] = "10"


@dataclass
class RMLVO:
    DEFAULT_RULES: ClassVar[str] = "evdev"
    DEFAULT_MODEL: ClassVar[str] = "pc105"
    DEFAULT_LAYOUT: ClassVar[str] = "us"
    rules: str
    model: str
    layout: str
    variant: str | None
    option: str | None

    @property
    def __iter__(self) -> Iterator[str | None]:
        yield self.rules
        yield self.model
        yield self.layout
        yield self.variant
        yield self.option

    @property
    def rmlvo(self) -> Iterator[tuple[str, str]]:
        yield ("rules", self.rules)
        yield ("model", self.model)
        yield ("layout", self.layout)
        # Keep only defined and non-empty values
        if self.variant is not None:
            yield ("variant", self.variant)
        if self.option is not None:
            yield ("option", self.option)

    @classmethod
    def from_rmlvo(
        cls,
        rules: str | None = None,
        model: str | None = None,
        layout: str | None = None,
        variant: str | None = None,
        option: str | None = None,
    ) -> Self:
        return cls(
            # We need to force a value for RML components
            rules or cls.DEFAULT_RULES,
            model or cls.DEFAULT_MODEL,
            layout or cls.DEFAULT_LAYOUT,
            variant,
            option,
        )


@dataclass
class Invocation(RMLVO, metaclass=ABCMeta):
    exitstatus: int = 77  # default to “skipped”
    error: str | None = None
    keymap: str = ""
    command: str = ""  # The fully compiled keymap

    def __str_iter(self) -> Iterator[str]:
        yield f"- rmlvo: {self.to_yaml(self.rmlvo)}"
        yield f'  cmd: "{self.escape(self.command)}"'
        yield f"  status: {self.exitstatus}"
        if self.error:
            yield f'  error: "{self.escape(self.error.strip())}"'

    def __str__(self) -> str:
        return "\n".join(self.__str_iter())

    @property
    def short(self) -> Iterator[tuple[str, str | int]]:
        yield from self.rmlvo
        yield ("status", self.exitstatus)
        if self.error is not None:
            yield ("error", self.error)

    @staticmethod
    def to_yaml(xs: Iterable[tuple[str, str | int]]) -> str:
        fields = ", ".join(f'{k}: "{v}"' for k, v in xs)
        return f"{{ {fields} }}"

    @staticmethod
    def escape(s: str) -> str:
        return s.replace('"', '\\"')

    def _write(self, fd: TextIO) -> None:
        fd.write(f"// {self.to_yaml(self.rmlvo)}\n")
        fd.write(self.keymap)

    def _write_keymap(self, output_dir: Path, compress: int) -> None:
        layout = self.layout
        if self.variant:
            layout += f"({self.variant})"
        (output_dir / self.model).mkdir(exist_ok=True)
        keymap_file = output_dir / self.model / layout
        if compress:
            keymap_file = keymap_file.with_suffix(".gz")
            with gzip.open(
                keymap_file, "wt", compresslevel=compress, encoding="utf-8"
            ) as fd:
                self._write(fd)
                fd.close()
        else:
            with keymap_file.open("wt", encoding="utf-8") as fd:
                self._write(fd)

    @classmethod
    @abstractmethod
    def run(cls, i: Self, output_dir: Path | None, compress: int) -> Self: ...

    @classmethod
    def run_all(
        cls,
        combos: Iterable[Self],
        combos_count: int,
        njobs: int,
        keymap_output_dir: Path | None,
        verbose: bool,
        short: bool,
        progress_bar: ProgressBar[Iterable[Self]],
        compress: int,
    ) -> bool:
        if keymap_output_dir:
            try:
                keymap_output_dir.mkdir(parents=True)
            except FileExistsError as e:
                print(e, file=sys.stderr)
                return False

        failed = False
        with multiprocessing.Pool(njobs) as p:
            f = partial(cls.run, output_dir=keymap_output_dir, compress=compress)
            results = p.imap_unordered(f, combos)
            for invocation in progress_bar(
                results, total=combos_count, file=sys.stdout
            ):
                if invocation.exitstatus != 0:
                    failed = True
                    target = sys.stderr
                else:
                    target = sys.stdout if verbose else None

                if target:
                    if short:
                        print("-", cls.to_yaml(invocation.short), file=target)
                    else:
                        print(invocation, file=target)

        return failed


@dataclass
class XkbCompInvocation(Invocation):
    xkbcomp_args: ClassVar[tuple[str, ...]] = ("xkbcomp", "-xkb", "-", "-")

    @classmethod
    def run(cls, i: Self, output_dir: Path | None, compress: int) -> Self:
        i._run(output_dir, compress)
        return i

    def _run(self, output_dir: Path | None, compress: int) -> None:
        args = (
            "setxkbmap",
            "-print",
            *itertools.chain.from_iterable((f"-{k}", v) for k, v in self.rmlvo),
        )

        self.command = " ".join(itertools.chain(args, "|", self.xkbcomp_args))

        setxkbmap = subprocess.Popen(
            args,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            universal_newlines=True,
        )
        stdout, stderr = setxkbmap.communicate()
        if "Cannot open display" in stderr:
            self.error = stderr
            self.exitstatus = 90
        else:
            xkbcomp = subprocess.Popen(
                self.xkbcomp_args,
                stdin=subprocess.PIPE,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                universal_newlines=True,
            )
            stdout, stderr = xkbcomp.communicate(stdout)
            if xkbcomp.returncode != 0:
                self.error = "failed to compile keymap"
                self.exitstatus = xkbcomp.returncode
            else:
                self.keymap = stdout
                self.exitstatus = 0

        if output_dir:
            self._write_keymap(output_dir, compress)


@dataclass
class XkbcommonInvocation(Invocation):
    UNRECOGNIZED_KEYSYM_ERROR: ClassVar[str] = "XKB-107"

    @classmethod
    def run(cls, i: Self, output_dir: Path | None, compress: int) -> Self:
        i._run(output_dir, compress)
        return i

    def _run(self, output_dir: Path | None, compress: int) -> None:
        args = (
            "xkbcli-compile-keymap",  # this is run in the builddir
            # Not needed, because we set XKB_LOG_LEVEL and XKB_LOG_VERBOSITY in env
            # "--verbose",
            *itertools.chain.from_iterable((f"--{k}", v) for k, v in self.rmlvo),
        )
        if not output_dir:
            args += ("--test",)
        self.command = " ".join(args)
        try:
            completed = subprocess.run(args, text=True, check=True, capture_output=True)
        except subprocess.CalledProcessError as err:
            self.error = "failed to compile keymap"
            self.exitstatus = err.returncode
        else:
            if self.UNRECOGNIZED_KEYSYM_ERROR in completed.stderr:
                for line in completed.stderr.splitlines():
                    if self.UNRECOGNIZED_KEYSYM_ERROR in line:
                        self.error = line
                        break
                self.exitstatus = 99  # tool doesn't generate this one
            else:
                self.exitstatus = 0
                self.keymap = completed.stdout

        if output_dir:
            self._write_keymap(output_dir, compress)


@dataclass
class Layout:
    name: str
    variants: list[str | None]

    @classmethod
    def parse(cls, e: ET.Element, variant: list[str] | None = None) -> Self:
        if (name_elem := e.find("configItem/name")) is None or name_elem is None:
            raise ValueError("Layout name not found")
        if not variant:
            variants = [None] + [
                cls.parse_text(v)
                for v in e.findall("variantList/variant/configItem/name")
            ]
        else:
            variants = cast(list[str | None], variant)
        return cls(cls.parse_text(e.find("configItem/name")), variants)

    @staticmethod
    def parse_text(e: ET.Element | None) -> str:
        if e is None or not e.text:
            raise ValueError("Name not found")
        return e.text


def parse_registry(
    paths: Sequence[Path],
    tool: type[Invocation],
    model: str | None,
    layout: str | None,
    variant: str | None,
    option: str | None,
) -> tuple[int, Iterator[Invocation]]:
    models: tuple[str, ...] = ()
    layouts: tuple[Layout, ...] = ()
    options: tuple[str, ...] = ()

    if variant and not layout:
        raise ValueError("Variant must be set together with layout")

    for path in paths:
        root = ET.fromstring(open(path).read())

        # Models
        if model is None:
            models += tuple(
                e.text
                for e in root.findall("modelList/model/configItem/name")
                if e.text
            )
        elif not models:
            models += (model,)

        # Layouts/variants
        if layout:
            if variant is None:
                layouts += tuple(
                    map(
                        Layout.parse,
                        (
                            e
                            for e in root.findall("layoutList/layout")
                            if e.find(f"configItem/name[.='{layout}']") is not None
                        ),
                    )
                )
            elif not layouts:
                layouts += (Layout(layout, cast(list[str | None], variant.split(":"))),)
        else:
            layouts += tuple(map(Layout.parse, root.findall("layoutList/layout")))

        # Options
        if option is None:
            options += tuple(
                e.text
                for e in root.findall("optionList/group/option/configItem/name")
                if e.text
            )
        elif not options and option:
            options += (option,)

    # Some registry may be only partial, e.g.: *.extras.xml
    if not models:
        models = (RMLVO.DEFAULT_MODEL,)
    if not layouts:
        layouts = (Layout(RMLVO.DEFAULT_LAYOUT, [None]),)

    count = len(models) * sum(len(l.variants) for l in layouts) * (1 + len(options))

    # The list of combos can be huge, so better to use a generator instead
    def iter_combos() -> Iterator[Invocation]:
        for m in models:
            for l in layouts:
                for v in l.variants:
                    yield tool.from_rmlvo(
                        rules=None, model=m, layout=l.name, variant=v, option=None
                    )
                    for opt in options:
                        yield tool.from_rmlvo(
                            rules=None, model=m, layout=l.name, variant=v, option=opt
                        )

    return count, iter_combos()


T = TypeVar("T")


# Needed because Callable does not handle keywords args
class ProgressBar(Protocol[T]):
    def __call__(self, x: T, total: int, file: TextIO | None) -> T: ...


# The function generating the progress bar (if any).
def create_progress_bar(verbose: bool) -> ProgressBar[T]:
    def noop_progress_bar(x: T, total: int, file: TextIO | None = None) -> T:
        return x

    progress_bar: ProgressBar[T] = noop_progress_bar
    if not verbose and os.isatty(sys.stdout.fileno()):
        try:
            from tqdm import tqdm

            progress_bar = cast(ProgressBar[T], tqdm)
        except ImportError:
            pass

    return progress_bar


def main() -> NoReturn:
    parser = argparse.ArgumentParser(
        description="""
                    This tool compiles a keymap for each layout, variant and
                    options combination in the given rules XML file. The output
                    of this tool is YAML, use your favorite YAML parser to
                    extract error messages. Errors are printed to stderr.
                    """
    )
    parser.add_argument(
        "paths",
        metavar="/path/to/evdev.xml",
        nargs="*",
        type=str,
        default=(DEFAULT_RULES_XML,),
        help="Path to xkeyboard-config's evdev.xml",
    )
    tools: dict[str, type[Invocation]] = {
        "libxkbcommon": XkbcommonInvocation,
        "xkbcomp": XkbCompInvocation,
    }
    parser.add_argument(
        "--tool",
        choices=tools.keys(),
        type=str,
        default="libxkbcommon",
        help="parsing tool to use",
    )
    parser.add_argument(
        "--jobs",
        "-j",
        type=int,
        default=4 * (os.cpu_count() or 1),
        help="number of processes to use",
    )
    parser.add_argument("--verbose", "-v", default=False, action="store_true")
    parser.add_argument(
        "--short", default=False, action="store_true", help="Concise output"
    )
    parser.add_argument(
        "--keymap-output-dir",
        default=None,
        type=Path,
        help="Directory to print compiled keymaps to",
    )
    parser.add_argument(
        "--compress", type=int, default=0, help="Compression level of keymaps files"
    )
    parser.add_argument(
        "--model", default="", type=str, help="Only test the given model"
    )
    parser.add_argument(
        "--layout", default=WILDCARD, type=str, help="Only test the given layout"
    )
    parser.add_argument(
        "--variant",
        default=WILDCARD,
        type=str,
        help="Only test the given variants (colon-separated list)",
    )
    parser.add_argument(
        "--option", default=WILDCARD, type=str, help="Only test the given option"
    )
    parser.add_argument(
        "--no-iterations", "-1", action="store_true", help="Only test one combo"
    )

    args = parser.parse_args()

    verbose: bool = args.verbose
    short = args.short
    keymapdir = args.keymap_output_dir
    progress_bar: ProgressBar[Iterable[Invocation]] = create_progress_bar(verbose)

    tool = tools[args.tool]

    model: str | None = None if args.model == WILDCARD else args.model
    layout: str | None = None if args.layout == WILDCARD else args.layout
    variant: str | None = None if args.variant == WILDCARD else args.variant
    option: str | None = None if args.option == WILDCARD else args.option

    if args.no_iterations:
        combos = (
            tool.from_rmlvo(
                rules=None, model=model, layout=layout, variant=variant, option=option
            ),
        )
        count = len(combos)
        iter_combos = iter(combos)
    else:
        count, iter_combos = parse_registry(
            args.paths, tool, model, layout, variant, option
        )

    failed = tool.run_all(
        iter_combos,
        count,
        args.jobs,
        keymapdir,
        verbose,
        short,
        progress_bar,
        args.compress,
    )
    sys.exit(failed)


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("# Exiting after Ctrl+C")