trezor-firmware/core/tools/build_mocks

#!/usr/bin/env python3
from __future__ import annotations

import os
import shutil
import subprocess
import sys
import tempfile
from pathlib import Path

import typing as t

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 = "///"
COMMENT_PREFIX_RS = "//#"

in_class = False
current_package = None
current_method_is_overload = False

def process_file(lines_iter: t.Iterator[str], block_prefix: str) -> t.Iterator[tuple[str, str]]:
    in_class = False
    package = None

    def process_block() -> t.Iterator[tuple[str, str]]:
        """Process a single block of comment-prefixed lines."""
        block_indent = "" if not in_class else " " * 4
        nonlocal package
        nonlocal in_class
        for line in lines_iter:
            line = line.strip()
            if not line.startswith(block_prefix):
                return
            
            line = line[len(block_prefix) + 1 :].rstrip()
            
            if line.startswith("package: "):
                package = line[len("package: ") :].strip()
                continue

            if package is None:
                raise ValueError("No package set (use 'package: <name>')")

        if not in_class:
            yield package, f"\n# {mod_desc}"

        for line in block:
            if not line.strip():
                continue

            yield package, line


def split_to_parts(line):
    global in_class
    global current_package
    global current_method_is_overload

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

    if current_package is None:
        raise ValueError("No package set (use 'package: <name>')")

    if line.startswith("mock:global"):
        in_class = False
        return

    indent = "" if not in_class else " " * 4
    yield (current_package, (indent + line).rstrip())

    if line.startswith("class "):
        in_class = True


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 + "\n")


def build_module(mod_file, dest):
    global current_indent
    global in_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
    in_class = None
    current_package = name.split("-")[0]
    mod_desc = str(mod_file.relative_to(CORE_DIR / "embed"))

    for l in open(mod_file):
        l = l.strip()
        if not l.startswith(COMMENT_PREFIX):
            continue

        l = l[len(COMMENT_PREFIX) + 1 :]  # account for space after ///
        store_to_file(dest, split_to_parts(l, mod_desc))


def build_rsmodule(mod_file):
    global current_indent
    global in_class
    global current_package

    assert mod_file.suffix == ".rs"

    current_indent = 0
    in_class = False
    current_package = None
    mod_desc = str(mod_file.relative_to(CORE_DIR / "embed"))

    yield_sep = False
    for l in open(mod_file):
        l = l.strip()
        if not l.startswith(COMMENT_PREFIX_RS):
            if yield_sep:
                if in_class:
                    yield current_package, ""
                else:
                    yield current_package, f"\n# {mod_desc}"
                yield_sep = False
            continue

        l = l[len(COMMENT_PREFIX_RS) + 1 :]  # account for space after //#
        try:
            yield from split_to_parts(l, mod_desc)
            yield_sep = True
        except ValueError as e:
            print(f"Error in {mod_file}: {e}")
            print(f"Line: {l}")
            sys.exit(1)


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")):
        store_to_file(dest, build_rsmodule(modfile))
    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()