setup.py

#!/usr/bin/env python
# flake8: noqa
import io
from contextlib import suppress, contextmanager
from os import fspath
from pathlib import Path
from typing import Optional, List, Dict

from setuptools import setup, Command, find_namespace_packages
from setuptools.command.build import build, SubCommand
from setuptools.command.editable_wheel import editable_wheel

import yaml


build.sub_commands.insert(0, ("compile-regexes", None))


class CompileRegexes(Command, SubCommand):
    def initialize_options(self) -> None:
        self.pkg_name: Optional[str] = None

    def finalize_options(self) -> None:
        self.pkg_name = self.distribution.get_name().replace("-", "_")

    def get_source_files(self) -> List[str]:
        return ["uap-core/regexes.yaml"]

    def get_outputs(self) -> List[str]:
        return [f"{self.pkg_name}/_regexes.py"]

    def get_output_mapping(self) -> Dict[str, str]:
        return dict(zip(self.get_source_files(), self.get_outputs()))

    def run(self) -> None:
        # FIXME: check git / submodules?
        """
        work_path = self.work_path
        if not os.path.exists(os.path.join(work_path, ".git")):
            return

        log.info("initializing git submodules")
        check_output(["git", "submodule", "init"], cwd=work_path)
        check_output(["git", "submodule", "update"], cwd=work_path)
        """
        if not self.pkg_name:
            return  # or error?

        yaml_src = Path("uap-core", "regexes.yaml")
        if not yaml_src.is_file():
            raise RuntimeError(
                f"Unable to find regexes.yaml, should be at {yaml_src!r}"
            )

        with yaml_src.open("rb") as f:
            regexes = yaml.safe_load(f)

        if self.editable_mode:
            dist_dir = Path("src")
        else:
            dist_dir = Path(self.get_finalized_command("bdist_wheel").bdist_dir)

        outdir = dist_dir / self.pkg_name
        outdir.mkdir(parents=True, exist_ok=True)

        dest = outdir / "_matchers.py"
        dest_lazy = outdir / "_lazy.py"
        dest_legacy = outdir / "_regexes.py"

        with (
            dest.open("wb") as eager,
            dest_lazy.open("wb") as lazy,
            dest_legacy.open("wb") as legacy,
        ):
            eager = EagerWriter(eager)
            lazy = LazyWriter(lazy)
            legacy = LegacyWriter(legacy)

            for section in ["user_agent_parsers", "os_parsers", "device_parsers"]:
                with (
                    eager.section(section),
                    lazy.section(section),
                    legacy.section(section),
                ):
                    extract = EXTRACTORS[section]
                    for p in regexes[section]:
                        el = trim(extract(p))
                        eager.item(el)
                        lazy.item(el)
                        legacy.item(el)
            eager.end()
            lazy.end()
            legacy.end()


def trim(l):
    while len(l) > 1 and l[-1] is None:
        l.pop()
    return l


EXTRACTORS = {
    "user_agent_parsers": lambda p: [
        p["regex"],
        p.get("family_replacement"),
        p.get("v1_replacement"),
        p.get("v2_replacement"),
    ],
    "os_parsers": lambda p: [
        p["regex"],
        p.get("os_replacement"),
        p.get("os_v1_replacement"),
        p.get("os_v2_replacement"),
        p.get("os_v3_replacement"),
        p.get("os_v4_replacement"),
    ],
    "device_parsers": lambda p: [
        p["regex"],
        p.get("regex_flag"),
        p.get("device_replacement"),
        p.get("brand_replacement"),
        p.get("model_replacement"),
    ],
}


class Writer:
    section_end = b""

    def __init__(self, fp):
        self.fp = fp
        self.fp.write(
            b"""\
########################################################
# NOTICE: this file is autogenerated from regexes.yaml #
########################################################
"""
        )
        self.fp.write(self.prefix)
        self._section = None

    @contextmanager
    def section(self, id):
        self._section = id
        self.fp.write(self.sections[id])
        yield
        self.fp.write(self.section_end)

    def item(self, elements):
        #        DeviceMatcher(re, flag, repl1),
        self.fp.write(self.items[self._section])
        self.fp.write(", ".join(map(repr, elements)).encode())
        self.fp.write(b"),\n")

    def end(self):
        self.fp.write(self.suffix)


class LegacyWriter(Writer):
    prefix = b"""\
__all__ = [
    "USER_AGENT_PARSERS",
    "DEVICE_PARSERS",
    "OS_PARSERS",
]

from .user_agent_parser import UserAgentParser, DeviceParser, OSParser

"""
    sections = {
        "user_agent_parsers": b"USER_AGENT_PARSERS = [\n",
        "os_parsers": b"\n\nOS_PARSERS = [\n",
        "device_parsers": b"\n\nDEVICE_PARSERS = [\n",
    }
    section_end = b"]"
    items = {
        "user_agent_parsers": b"    UserAgentParser(",
        "os_parsers": b"    OSParser(",
        "device_parsers": b"    DeviceParser(",
    }
    suffix = b"\n"


class EagerWriter(Writer):
    prefix = b"""\
__all__ = ["MATCHERS"]

from typing import Tuple, List
from .matchers import UserAgentMatcher, OSMatcher, DeviceMatcher

MATCHERS: Tuple[List[UserAgentMatcher], List[OSMatcher], List[DeviceMatcher]] = ([
"""
    sections = {
        "user_agent_parsers": b"",
        "os_parsers": b"], [\n",
        "device_parsers": b"], [\n",
    }
    items = {
        "user_agent_parsers": b"    UserAgentMatcher(",
        "os_parsers": b"    OSMatcher(",
        "device_parsers": b"    DeviceMatcher(",
    }
    suffix = b"])\n"


class LazyWriter(EagerWriter):
    prefix = b"""\
__all__ = ["MATCHERS"]

from typing import Tuple, List
from .lazy import UserAgentMatcher, OSMatcher, DeviceMatcher

MATCHERS: Tuple[List[UserAgentMatcher], List[OSMatcher], List[DeviceMatcher]] = ([
"""


setup(
    cmdclass={
        "compile-regexes": CompileRegexes,
    }
)