proto_prune_tree.py

#!/usr/bin/env python3
"""
proto_prune_tree.py

Copy only the .proto files and declarations needed for one protobuf entry into a
new folder. File names and relative directories are preserved, but each copied
file is stripped to its required messages/enums/services.

Entries:
  file.proto:RootlistModificationRequest
  file.proto:RootlistModificationService.Modify

For a service-method entry, the output keeps only that RPC method in the service
and recursively keeps its request and response types.

Examples:
  python proto_prune_tree.py \
    --src ./protos \
    --entry rootlist_modification_request.proto:RootlistModificationRequest \
    --dst ./minimal-protos

  python proto_prune_tree.py \
    --src ./protos \
    --entry rootlist_modification_service.proto:RootlistModificationService.Modify \
    --dst ./minimal-protos
"""

from __future__ import annotations

import argparse
import json
import re
from collections import defaultdict, deque
from dataclasses import dataclass, field
from pathlib import Path


SCALAR_TYPES = {
    "double", "float",
    "int32", "int64", "uint32", "uint64", "sint32", "sint64",
    "fixed32", "fixed64", "sfixed32", "sfixed64",
    "bool", "string", "bytes",
}

DECL_RE = re.compile(r"\b(message|enum|service)\s+([A-Za-z_]\w*)\b")
PACKAGE_RE = re.compile(r"(?m)^\s*package\s+([A-Za-z_]\w*(?:\.[A-Za-z_]\w*)*)\s*;")
SYNTAX_RE = re.compile(r'(?m)^\s*syntax\s*=\s*"([^"]+)"\s*;')
IMPORT_RE = re.compile(r'(?m)^\s*import\s+(?:(public|weak)\s+)?"([^"]+)"\s*;')
# File-level options such as java_package, java_multiple_files, objc_class_prefix,
# and custom extension options. They are preserved verbatim in pruned output.
OPTION_RE = re.compile(r"(?m)^\s*option\s+[^;]+;")

FIELD_TYPE_RE = re.compile(
    r"""(?:
            \b(?:optional|required|repeated)\s+
        )?
        (?:
            map\s*<\s*([.\w]+)\s*,\s*([.\w]+)\s*>
            |
            ([.\w]+)
        )
        \s+[A-Za-z_]\w*\s*=""",
    re.X,
)

RPC_RE = re.compile(
    r"\brpc\s+([A-Za-z_]\w*)\s*"
    r"\(\s*(?:stream\s+)?([.\w]+)\s*\)\s*"
    r"returns\s*\(\s*(?:stream\s+)?([.\w]+)\s*\)"
    r"\s*(?:\{.*?\}|;)",
    re.S,
)


@dataclass(frozen=True)
class ImportSpec:
    qualifier: str | None
    path: str


@dataclass
class RpcMethod:
    name: str
    request_type: str
    response_type: str
    text: str


@dataclass
class Declaration:
    kind: str
    name: str
    package: str
    source_rel: str
    text: str
    refs: set[str] = field(default_factory=set)
    rpc_methods: dict[str, RpcMethod] = field(default_factory=dict)
    selected_rpc_methods: set[str] | None = None

    @property
    def fqn(self) -> str:
        return f"{self.package}.{self.name}" if self.package else self.name


@dataclass
class ProtoFile:
    source_rel: str
    syntax: str
    package: str
    imports: list[ImportSpec]
    options: list[str]
    declarations: list[Declaration]


def strip_comments(source: str) -> str:
    """Remove comments without changing string literals or newlines."""
    out: list[str] = []
    index = 0
    quote: str | None = None

    while index < len(source):
        char = source[index]
        next_char = source[index + 1] if index + 1 < len(source) else ""

        if quote:
            out.append(char)
            if char == "\\" and index + 1 < len(source):
                index += 1
                out.append(source[index])
            elif char == quote:
                quote = None
            index += 1
            continue

        if char in ("'", '"'):
            quote = char
            out.append(char)
            index += 1
        elif char == "/" and next_char == "/":
            while index < len(source) and source[index] != "\n":
                out.append(" ")
                index += 1
        elif char == "/" and next_char == "*":
            out.extend((" ", " "))
            index += 2
            while index < len(source) - 1:
                if source[index] == "*" and source[index + 1] == "/":
                    out.extend((" ", " "))
                    index += 2
                    break
                out.append("\n" if source[index] == "\n" else " ")
                index += 1
        else:
            out.append(char)
            index += 1

    return "".join(out)


def matching_brace(source: str, opening_brace: int) -> int:
    depth = 0
    quote: str | None = None
    escaped = False

    for index in range(opening_brace, len(source)):
        char = source[index]

        if quote:
            if escaped:
                escaped = False
            elif char == "\\":
                escaped = True
            elif char == quote:
                quote = None
            continue

        if char in ("'", '"'):
            quote = char
        elif char == "{":
            depth += 1
        elif char == "}":
            depth -= 1
            if depth == 0:
                return index

    raise ValueError(f"Unclosed declaration brace at offset {opening_brace}")


def extract_refs(declaration_text: str) -> set[str]:
    clean = strip_comments(declaration_text)
    result: set[str] = set()

    for match in FIELD_TYPE_RE.finditer(clean):
        for type_name in match.groups():
            if type_name:
                type_name = type_name.lstrip(".")
                if type_name not in SCALAR_TYPES:
                    result.add(type_name)

    for match in RPC_RE.finditer(clean):
        request_type = match.group(2).lstrip(".")
        response_type = match.group(3).lstrip(".")
        if request_type not in SCALAR_TYPES:
            result.add(request_type)
        if response_type not in SCALAR_TYPES:
            result.add(response_type)

    return result


def parse_rpc_methods(service_text: str) -> dict[str, RpcMethod]:
    clean = strip_comments(service_text)
    methods: dict[str, RpcMethod] = {}

    for match in RPC_RE.finditer(clean):
        name = match.group(1)
        methods[name] = RpcMethod(
            name=name,
            request_type=match.group(2).lstrip("."),
            response_type=match.group(3).lstrip("."),
            # Match offsets are valid because comments are replaced by equal length.
            text=service_text[match.start():match.end()].strip(),
        )

    return methods



JAVA_PACKAGE_OPTION_RE = re.compile(
    r'^(?P<prefix>\s*option\s+java_package\s*=\s*")(?P<package>[^"]+)(?P<suffix>"\s*;)\s*$'
)

def normalize_java_package_option(option_text: str) -> str:
    """
    Prefix java_package values with 'com.' only when they start exactly with
    'spotify.'. Values already beginning with 'com.spotify.' stay unchanged.
    """
    match = JAVA_PACKAGE_OPTION_RE.match(option_text)
    if not match:
        return option_text

    package_name = match.group("package")
    if package_name == "spotify" or package_name.startswith("spotify."):
        package_name = f"com.{package_name}"

    return (
        f'{match.group("prefix")}{package_name}{match.group("suffix")}'
    )


def parse_proto(path: Path, src_root: Path) -> ProtoFile:
    raw = path.read_text(encoding="utf-8")
    clean = strip_comments(raw)
    source_rel = path.relative_to(src_root).as_posix()

    package_match = PACKAGE_RE.search(clean)
    syntax_match = SYNTAX_RE.search(clean)
    package = package_match.group(1) if package_match else ""
    syntax = syntax_match.group(1) if syntax_match else "proto3"

    imports = [
        ImportSpec(match.group(1), match.group(2))
        for match in IMPORT_RE.finditer(clean)
    ]

    # Use the raw source slices so formatting/comments in options survive.
    options = [
        normalize_java_package_option(raw[match.start():match.end()].strip())
        for match in OPTION_RE.finditer(clean)
    ]

    declarations: list[Declaration] = []
    occupied_until = -1

    for match in DECL_RE.finditer(clean):
        if match.start() < occupied_until:
            continue

        opening_brace = clean.find("{", match.end())
        if opening_brace < 0:
            continue

        if clean.find(";", match.end(), opening_brace) != -1:
            continue

        closing_brace = matching_brace(clean, opening_brace)
        occupied_until = closing_brace + 1

        kind, name = match.group(1), match.group(2)
        text = raw[match.start():closing_brace + 1].strip()
        decl = Declaration(
            kind=kind,
            name=name,
            package=package,
            source_rel=source_rel,
            text=text,
            refs=extract_refs(text),
        )
        if kind == "service":
            decl.rpc_methods = parse_rpc_methods(text)
        declarations.append(decl)

    return ProtoFile(source_rel, syntax, package, imports, options, declarations)


def resolve_ref(
        raw_ref: str,
        owner: Declaration,
        by_fqn: dict[str, Declaration],
        by_simple_name: dict[str, list[Declaration]],
) -> Declaration | None:
    ref = raw_ref.lstrip(".")
    if ref in SCALAR_TYPES:
        return None

    if ref in by_fqn:
        return by_fqn[ref]

    # Nested type references are retained via the surrounding top-level declaration.
    segments = ref.split(".")
    for end in range(len(segments), 0, -1):
        candidate = ".".join(segments[:end])
        if candidate in by_fqn:
            return by_fqn[candidate]

    if "." not in ref and owner.package:
        candidate = f"{owner.package}.{ref}"
        if candidate in by_fqn:
            return by_fqn[candidate]

    candidates = by_simple_name.get(ref, [])
    return candidates[0] if len(candidates) == 1 else None


def parse_entry(value: str) -> tuple[str, str, str | None]:
    """
    Parse:
      foo.proto:Message
      foo.proto:Service.Method
    """
    if ":" not in value:
        raise argparse.ArgumentTypeError(
            "Entry must be relative/file.proto:Declaration or relative/file.proto:Service.Method"
        )

    source_file, target = value.rsplit(":", 1)
    parts = target.split(".")

    if len(parts) == 1:
        return source_file.replace("\\", "/"), parts[0], None
    if len(parts) == 2:
        return source_file.replace("\\", "/"), parts[0], parts[1]

    raise argparse.ArgumentTypeError(
        "Only top-level Service.Method entries are supported."
    )


def render_service_subset(decl: Declaration) -> str:
    """Keep only selected RPC methods when the entry was Service.Method."""
    if not decl.selected_rpc_methods:
        return decl.text

    methods = [
        decl.rpc_methods[name].text
        for name in sorted(decl.selected_rpc_methods)
    ]

    package_prefix = ""
    match = re.match(r"(service\s+[A-Za-z_]\w*\s*)\{", decl.text, re.S)
    if not match:
        return decl.text

    package_prefix = match.group(1)
    body = "\n\n".join(
        "  " + method.replace("\n", "\n  ")
        for method in methods
    )
    return f"{package_prefix}{{\n{body}\n}}"


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Create a dependency-minimal, stripped .proto file tree."
    )
    parser.add_argument("--src", required=True, type=Path,
                        help="Root directory containing original .proto files.")
    parser.add_argument("--entry", required=True, type=parse_entry,
                        help="file.proto:Declaration or file.proto:Service.Method")
    parser.add_argument("--dst", required=True, type=Path,
                        help="New output directory for pruned .proto files.")
    # Existing destinations are merged into; existing declarations/methods are preserved.
    args = parser.parse_args()

    src = args.src.resolve()
    dst = args.dst.resolve()
    entry_file, entry_declaration_name, entry_method = args.entry

    if not src.is_dir():
        parser.error(f"Source directory does not exist: {src}")

    # Do not delete an existing destination. Existing stripped files are parsed and
    # merged with the new dependency closure below.
    source_files = sorted(src.rglob("*.proto"))
    if not source_files:
        parser.error(f"No .proto files found under: {src}")

    protos = {
        path.relative_to(src).as_posix(): parse_proto(path, src)
        for path in source_files
    }

    if entry_file not in protos:
        parser.error(f"Entry file does not exist beneath --src: {entry_file}")

    candidates = [
        decl for decl in protos[entry_file].declarations
        if decl.name == entry_declaration_name
    ]
    if not candidates:
        parser.error(
            f"Declaration {entry_declaration_name!r} not found in {entry_file}"
        )

    entry = candidates[0]

    # Explicit method existence validation requested by the user.
    if entry_method is not None:
        if entry.kind != "service":
            parser.error(
                f"{entry_declaration_name} is a {entry.kind}, not a service; "
                f"cannot select method {entry_method!r}."
            )
        if entry_method not in entry.rpc_methods:
            available = ", ".join(sorted(entry.rpc_methods)) or "(none)"
            parser.error(
                f"Method {entry_method!r} does not exist in service "
                f"{entry.fqn}. Available methods: {available}"
            )
        entry.selected_rpc_methods = {entry_method}

    by_fqn: dict[str, Declaration] = {}
    by_simple_name: dict[str, list[Declaration]] = defaultdict(list)
    for proto in protos.values():
        for decl in proto.declarations:
            by_fqn[decl.fqn] = decl
            by_simple_name[decl.name].append(decl)

    selected: dict[str, Declaration] = {}
    unresolved: dict[str, set[str]] = defaultdict(set)
    pending: deque[Declaration] = deque([entry])

    while pending:
        decl = pending.popleft()
        existing = selected.get(decl.fqn)

        # If this is a selected method from a service, merge method selection.
        if existing is not None:
            if decl.selected_rpc_methods:
                existing.selected_rpc_methods = (
                        (existing.selected_rpc_methods or set())
                        | decl.selected_rpc_methods
                )
            continue

        selected[decl.fqn] = decl

        refs = set(decl.refs)
        if decl.kind == "service" and decl.selected_rpc_methods:
            refs = set()
            for method_name in decl.selected_rpc_methods:
                method = decl.rpc_methods[method_name]
                refs.add(method.request_type)
                refs.add(method.response_type)

        for ref in sorted(refs):
            dependency = resolve_ref(ref, decl, by_fqn, by_simple_name)
            if dependency is None:
                unresolved[decl.fqn].add(ref)
            elif dependency.fqn not in selected:
                pending.append(dependency)

    selected_by_file: dict[str, list[Declaration]] = defaultdict(list)
    for decl in selected.values():
        selected_by_file[decl.source_rel].append(decl)

    selected_fqns = set(selected)

    dst.mkdir(parents=True, exist_ok=True)
    added_declarations: list[str] = []
    skipped_declarations: list[str] = []
    added_methods: list[str] = []
    skipped_methods: list[str] = []

    for source_rel, declarations in sorted(selected_by_file.items()):
        proto = protos[source_rel]
        output_path = dst / source_rel
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # Parse the existing generated/pruned file, when present. This lets repeated
        # invocations accumulate declarations and service methods instead of replacing
        # or duplicating them.
        existing_proto = None
        existing_by_name: dict[str, Declaration] = {}
        if output_path.exists():
            existing_proto = parse_proto(output_path, dst)
            existing_by_name = {
                declaration.name: declaration
                for declaration in existing_proto.declarations
            }

        declarations_to_append: list[Declaration] = []
        existing_service_methods: dict[str, set[str]] = {
            name: set(decl.rpc_methods)
            for name, decl in existing_by_name.items()
            if decl.kind == "service"
        }

        for decl in sorted(declarations, key=lambda item: item.name):
            old_decl = existing_by_name.get(decl.name)

            if old_decl is None:
                declarations_to_append.append(decl)
                added_declarations.append(decl.fqn)
                if decl.kind == "service" and decl.selected_rpc_methods:
                    added_methods.extend(
                        f"{decl.fqn}.{method}"
                        for method in sorted(decl.selected_rpc_methods)
                    )
                continue

            # Existing non-service declaration: exact name already exists, so retain it.
            if decl.kind != "service":
                skipped_declarations.append(decl.fqn)
                continue

            # Existing service declaration: only append the RPC methods that were not
            # already written to the stripped output.
            selected_methods = decl.selected_rpc_methods
            if selected_methods is None:
                # A full service was requested. Keep only methods not already present.
                selected_methods = set(decl.rpc_methods)

            missing_methods = selected_methods - existing_service_methods.get(decl.name, set())
            duplicate_methods = selected_methods - missing_methods

            skipped_methods.extend(
                f"{decl.fqn}.{method}"
                for method in sorted(duplicate_methods)
            )

            if missing_methods:
                # Append a small reopening of the existing service rather than rewrite
                # the file. Protobuf allows a service only once, so instead we mark the
                # file for a complete safe rewrite below.
                clone = Declaration(
                    kind=decl.kind,
                    name=decl.name,
                    package=decl.package,
                    source_rel=decl.source_rel,
                    text=decl.text,
                    refs=decl.refs,
                    rpc_methods=decl.rpc_methods,
                    selected_rpc_methods=missing_methods,
                )
                declarations_to_append.append(clone)
                added_methods.extend(
                    f"{decl.fqn}.{method}"
                    for method in sorted(missing_methods)
                )
            else:
                skipped_declarations.append(decl.fqn)

        # Imports needed by this invocation's selected closure.
        needed_source_files: set[str] = set()
        has_unresolved = False
        for decl in declarations:
            refs = set(decl.refs)
            if decl.kind == "service" and decl.selected_rpc_methods:
                refs = set()
                for method_name in decl.selected_rpc_methods:
                    method = decl.rpc_methods[method_name]
                    refs.update((method.request_type, method.response_type))

            for ref in refs:
                dependency = resolve_ref(ref, decl, by_fqn, by_simple_name)
                if dependency and dependency.fqn in selected_fqns:
                    if dependency.source_rel != source_rel:
                        needed_source_files.add(dependency.source_rel)
                elif dependency is None:
                    has_unresolved = True

        requested_import_lines: list[str] = []
        for imp in proto.imports:
            if imp.path in needed_source_files or has_unresolved:
                qualifier = f"{imp.qualifier} " if imp.qualifier else ""
                requested_import_lines.append(f'import {qualifier}"{imp.path}";')

        if existing_proto is None:
            # Fresh file: render the normal stripped proto.
            chunks = [f'syntax = "{proto.syntax}";', ""]
            if proto.package:
                chunks.extend((f"package {proto.package};", ""))
            if proto.options:
                chunks.extend(proto.options)
                chunks.append("")
            if requested_import_lines:
                chunks.extend(requested_import_lines)
                chunks.append("")
            for index, decl in enumerate(sorted(declarations_to_append, key=lambda item: item.name)):
                chunks.append(render_service_subset(decl))
                if index != len(declarations_to_append) - 1:
                    chunks.append("")
            output_path.write_text("\n".join(chunks) + "\n", encoding="utf-8")
            continue

        # Existing file:
        # - Add only missing imports.
        # - For an existing service, rebuild the one service declaration with the
        #   union of prior and newly requested methods, so the proto stays valid.
        existing_text = output_path.read_text(encoding="utf-8")

        # Normalize an existing unqualified Spotify java_package as well, so a
        # rerun upgrades prior output without adding a duplicate option.
        existing_text = "\n".join(
            normalize_java_package_option(line)
            if line.strip().startswith("option java_package")
            else line
            for line in existing_text.splitlines()
        ) + ("\n" if existing_text.endswith("\n") else "")

        existing_imports = {
            line.strip()
            for line in existing_text.splitlines()
            if line.strip().startswith("import ")
        }
        missing_import_lines = [
            line for line in requested_import_lines
            if line not in existing_imports
        ]

        rebuilt_services: dict[str, Declaration] = {}
        plain_additions: list[Declaration] = []

        for decl in declarations_to_append:
            if decl.kind != "service" or decl.name not in existing_by_name:
                plain_additions.append(decl)
                continue

            original_service = next(
                source_decl for source_decl in proto.declarations
                if source_decl.name == decl.name and source_decl.kind == "service"
            )
            union_methods = (
                    set(existing_by_name[decl.name].rpc_methods)
                    | set(decl.selected_rpc_methods or ())
            )
            original_service.selected_rpc_methods = union_methods
            rebuilt_services[decl.name] = original_service

        # Replace an existing service block in place with its method union.
        for service_name, replacement in rebuilt_services.items():
            old_service = existing_by_name[service_name]
            old_text = old_service.text
            replacement_text = render_service_subset(replacement)
            existing_text = existing_text.replace(old_text, replacement_text, 1)

        # Add file-level options that were not already retained.
        existing_option_lines = {
            line.strip()
            for line in existing_text.splitlines()
            if line.strip().startswith("option ")
        }
        missing_option_lines = [
            option for option in proto.options
            if option.strip() not in existing_option_lines
        ]
        if missing_option_lines:
            package_match = re.search(r'(?m)^\s*package\s+[^;]+;\s*$', existing_text)
            syntax_match = re.search(r'(?m)^\s*syntax\s*=\s*"[^"]+"\s*;\s*$', existing_text)
            insertion_point = package_match.end() if package_match else syntax_match.end()
            existing_text = (
                    existing_text[:insertion_point]
                    + "\n\n"
                    + "\n".join(missing_option_lines)
                    + existing_text[insertion_point:]
            )

        # Add imports after package/options (or syntax if no package).
        if missing_import_lines:
            option_matches = list(OPTION_RE.finditer(existing_text))
            package_match = re.search(r'(?m)^\s*package\s+[^;]+;\s*$', existing_text)
            syntax_match = re.search(r'(?m)^\s*syntax\s*=\s*"[^"]+"\s*;\s*$', existing_text)
            insertion_point = (
                option_matches[-1].end()
                if option_matches
                else (package_match.end() if package_match else syntax_match.end())
            )
            existing_text = (
                    existing_text[:insertion_point]
                    + "\n\n"
                    + "\n".join(missing_import_lines)
                    + existing_text[insertion_point:]
            )

        # Append previously absent messages/enums/new services.
        if plain_additions:
            existing_text = existing_text.rstrip() + "\n\n"
            existing_text += "\n\n".join(
                render_service_subset(decl)
                for decl in plain_additions
            ) + "\n"

        output_path.write_text(existing_text, encoding="utf-8")

    manifest = {
        "entry": {
            "source_file": entry.source_rel,
            "declaration": entry.fqn,
            "method": entry_method,
        },
        "files_written": sorted(selected_by_file),
        "declarations_written": sorted(selected),
        "unresolved_references": {
            fqn: sorted(refs)
            for fqn, refs in sorted(unresolved.items())
            if refs
        },
        "merge": {
            "added_declarations": sorted(set(added_declarations)),
            "skipped_existing_declarations": sorted(set(skipped_declarations)),
            "added_methods": sorted(set(added_methods)),
            "skipped_existing_methods": sorted(set(skipped_methods)),
        },
    }
    dst.mkdir(parents=True, exist_ok=True)
    (dst / "proto_prune_manifest.json").write_text(
        json.dumps(manifest, indent=2) + "\n",
        encoding="utf-8",
        )

    print(f"Entry:        {entry.fqn}" + (f".{entry_method}" if entry_method else ""))
    print(f"Files updated: {len(selected_by_file)}")
    print(f"Declarations in closure: {len(selected)}")
    print(f"Added declarations: {len(set(added_declarations))}")
    print(f"Added methods: {len(set(added_methods))}")
    print(f"Skipped existing methods: {len(set(skipped_methods))}")
    print(f"Output: {dst}")
    if unresolved:
        print("Warning: unresolved references are listed in proto_prune_manifest.json")
        return 2
    return 0


if __name__ == "__main__":
    raise SystemExit(main())