Materialize/ci/test/lint-buf/generate-buf-config.py

#!/usr/bin/env python3

# Copyright Materialize, Inc. and contributors. All rights reserved.
#
# Use of this software is governed by the Business Source License
# included in the LICENSE file at the root of this repository.
#
# As of the Change Date specified in that file, in accordance with
# the Business Source License, use of this software will be governed
# by the Apache License, Version 2.0.

import glob
import os
import re

BUF_YAML_FILE_PATH = "./src/buf.yaml"
BUF_YAML_TEMPLATE_FILE_PATH = f"{BUF_YAML_FILE_PATH}.template"

SOURCE_DIR = "src/"
PROTO_FILE_GLOB = f"{SOURCE_DIR}**/*.proto"

GENERATION_COMMENT = "File generated by generate-buf-config.py - DO NOT EDIT"
BUF_INSTRUCTION_PREFIX = "// buf breaking:"
# matched examples: "ignore", "ignore (database-issues#1000 must be fixed)"
BUF_IGNORE_COMMAND = re.compile(r"ignore( \((.*)\))?")


class ProtoFile:
    def __init__(self, path: str):
        self.path = path
        self.ignore_reason: str | None = None

    def is_ignore(self) -> bool:
        return self.ignore_reason is not None


def collect_proto_files() -> list[ProtoFile]:
    print(f"Working dir: {os.getcwd()}")
    proto_file_paths = glob.glob(PROTO_FILE_GLOB, recursive=True)
    return [ProtoFile(path) for path in proto_file_paths]


def load_buf_instructions(files: list[ProtoFile]) -> None:
    for file in files:
        load_buf_instructions_for_file(file)


def load_buf_instructions_for_file(file: ProtoFile) -> None:
    with open(file.path) as lines:
        for line in lines:
            if line.startswith(BUF_INSTRUCTION_PREFIX):
                handle_buf_instruction_in_proto_file(file, line)


def handle_buf_instruction_in_proto_file(file: ProtoFile, line: str) -> None:
    command = line.removeprefix(BUF_INSTRUCTION_PREFIX).strip()
    match = BUF_IGNORE_COMMAND.search(command)

    if match:
        has_ignore_reason = len(match.groups()) == 2 and match.group(2)
        # groups are not zero-based because group 0 contains the whole match
        file.ignore_reason = (
            match.group(2) if has_ignore_reason else "no reason specified"
        )
    else:
        raise RuntimeError(f"Unsupported buf instruction in {file.path}: {line}")


def generate_buf_ignore_section(ignored_files: list[ProtoFile]) -> str:
    ignore_entry_lines = []

    ignored_files = sorted(ignored_files, key=lambda file: file.path)

    for ignored_file in ignored_files:
        relative_path = ignored_file.path.removeprefix(SOURCE_DIR)
        ignore_entry_lines.append(f"    # reason: {ignored_file.ignore_reason}")
        ignore_entry_lines.append(f"    - {relative_path}")

    if len(ignore_entry_lines) == 0:
        ignore_entry_lines.append("     # none")

    return "\n".join(ignore_entry_lines).strip()


def write_buf_configuration(
    template_path: str, target_path: str, ignored_files: list[ProtoFile]
) -> None:

    with open(template_path) as input_file:
        content = input_file.read()
        content = content.replace("${generation-comment}", GENERATION_COMMENT)
        content = content.replace(
            "${ignore-entries}", generate_buf_ignore_section(ignored_files)
        )

    with open(target_path, "w") as output_file:
        output_file.write(content)


def main() -> None:
    proto_files = collect_proto_files()
    print(f"Collected {len(proto_files)} proto files.")

    load_buf_instructions(proto_files)
    ignored_proto_files = [file for file in proto_files if file.is_ignore()]
    print(f"{len(ignored_proto_files)} proto files to be ignored from breaking check.")

    write_buf_configuration(
        BUF_YAML_TEMPLATE_FILE_PATH, BUF_YAML_FILE_PATH, ignored_proto_files
    )
    print(f"Written buf configuration to '{BUF_YAML_FILE_PATH}'.")


if __name__ == "__main__":
    main()