trezor-firmware/core/tools/build_mocks

#!/usr/bin/env python3
import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

CORE_DIR = Path(__file__).resolve().parent.parent
EXTMOD_PATH = CORE_DIR / "embed" / "extmod"
RUSTMOD_PATH = CORE_DIR / "embed" / "rust" / "src"
MOCKS_PATH = CORE_DIR / "mocks" / "generated"

COMMENT_PREFIX_SHORT = "///"
COMMENT_PREFIX = COMMENT_PREFIX_SHORT + " "

current_indent = 0
current_class = None
current_method = None
current_package = None
current_method_is_overload = False


def split_to_parts(line, mod_desc=None):
    global current_indent
    global current_class
    global current_method
    global current_package
    global current_method_is_overload

    if line.startswith("package: "):
        current_package = line[9:].strip()
        return

    if line.startswith("mock:global"):
        current_indent = 0
        current_class = None
        return

    if line == "@overload\n":
        current_method_is_overload = True
        return

    if line.startswith("class "):
        current_class = line[6:].split("(")[0].strip(":")
        current_indent = 0

        yield (current_package, "\n\n")
        yield (current_package, "# " + mod_desc + "\n")

    elif line.startswith("def "):
        current_method = line[4:].split("(")[0]

        if current_class is None:
            yield (current_package, "\n\n")
            yield (current_package, "# " + mod_desc + "\n")
        else:
            yield (current_package, "\n")
            current_indent = 4

        if current_method_is_overload:
            yield (current_package, current_indent * " " + "@overload\n")
            current_method_is_overload = False

    if not line.strip():
        yield (current_package, "\n")
    else:
        line = current_indent * " " + line
        yield (current_package, line)


def store_to_file(dest, parts):
    for package, line in parts:
        package = package.replace(".", "/")
        dirpath = dest / os.path.dirname(package)
        filename = os.path.basename(package) + ".pyi"
        filepath = dirpath / filename

        dirpath.mkdir(parents=True, exist_ok=True)

        if (dest / package).is_dir():
            if not line.strip():
                continue
            print(f"Package exists: {package}")
            print(f"You should set 'package:' in {line.strip()}")
            sys.exit(1)

        if not filepath.exists():
            filepath.write_text("from typing import *\n")

        with open(filepath, "a") as f:
            f.write(line)


def build_module(mod_file, dest):
    global current_indent
    global current_class
    global current_package

    assert mod_file.name.startswith("mod")
    assert mod_file.suffix in (".c", ".h")
    # modfoobar-xyz.h -> foobar-xyz
    name = mod_file.name[3:-2]

    current_indent = 0
    current_class = None
    current_package = name.split("-")[0]
    mod_desc = str(mod_file.relative_to(CORE_DIR / "embed"))

    for l in open(mod_file):
        if not l.startswith(COMMENT_PREFIX_SHORT):
            continue
        l = l[len(COMMENT_PREFIX) :]  # .strip()

        store_to_file(dest, split_to_parts(l, mod_desc))


def build_rsmodule(mod_file, dest):
    global current_indent
    global current_class
    global current_package

    assert mod_file.suffix == ".rs"

    start_prefix = "pub static mp_module_"
    comment_prefix = f"    {COMMENT_PREFIX}"
    in_module = False
    current_indent = 0
    current_class = None
    mod_desc = str(mod_file.relative_to(CORE_DIR / "embed"))

    for l in open(mod_file):
        if l.startswith(start_prefix):
            in_module = True
            current_package = l[len(start_prefix) : ].split(":")[0]
        elif l.startswith("}"):
            in_module = False

        if not in_module:
            continue

        if not l.startswith(comment_prefix):
            continue

        l = l[len(comment_prefix) :]
        store_to_file(dest, split_to_parts(l, mod_desc))


def place_symlinks(dest):
    # make symlinks for the non-generated files
    for pyi in MOCKS_PATH.glob("../*.pyi"):
        dest_file = dest / pyi.name
        dest_file.symlink_to(os.path.relpath(pyi.resolve(), dest))


def build_directory(dest):
    for modfile in sorted(EXTMOD_PATH.glob("**/mod*.[ch]")):
        build_module(modfile, dest)
    for modfile in sorted(RUSTMOD_PATH.glob("**/*.rs")):
        build_rsmodule(modfile, dest)
    place_symlinks(dest)


def do_check():
    with tempfile.TemporaryDirectory() as tmpdir:
        build_directory(Path(tmpdir))
        diff_out = subprocess.run(
            ["diff", "-ur", str(MOCKS_PATH), tmpdir],
            stdout=subprocess.PIPE,
            universal_newlines=True,
        ).stdout
        if diff_out.strip():
            print(diff_out, end="")
            sys.exit(1)


def do_generate():
    shutil.rmtree(MOCKS_PATH)
    build_directory(MOCKS_PATH)


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--check":
        do_check()
    else:
        do_generate()