wizard-codegen 0.1.5__py3-none-any.whl

cli/__init__.py

@@ -0,0 +1,28 @@
+# Re-export everything from main for backwards compatibility
+from cli.main import (
+    app,
+    console,
+    Ctx,
+    main,
+    common,
+    generate,
+    list_protos,
+    validate,
+    _print_files_table,
+    _print_verbose_enabled,
+    _load_config,
+)
+
+__all__ = [
+    "app",
+    "console",
+    "Ctx",
+    "main",
+    "common",
+    "generate",
+    "list_protos",
+    "validate",
+    "_print_files_table",
+    "_print_verbose_enabled",
+    "_load_config",
+]

cli/main.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+import typer
+from pathlib import Path
+from rich.console import Console
+from rich.table import Table
+from rich.panel import Panel
+from rich.pretty import Pretty
+from dataclasses import dataclass
+from importlib.metadata import version as get_version
+from typing import Optional
+
+from core import *
+from proto import *
+import shutil
+
+app = typer.Typer(no_args_is_help=True, rich_markup_mode="rich")
+console = Console()
+
+
+def version_callback(value: bool):
+    if value:
+        ver = get_version("wizard-codegen")
+        console.print(f"wizard-codegen [bold cyan]{ver}[/]")
+        raise typer.Exit()
+
+
+@dataclass
+class Ctx:
+    verbose: bool = False
+    dry_run: bool = False
+    local: bool = False
+    config_path: Path = Path("wizard/codegen.yaml")
+
+
+@app.callback()
+def common(
+        ctx: typer.Context,
+        version: Optional[bool] = typer.Option(
+            None, "--version", "-V",
+            callback=version_callback,
+            is_eager=True,
+            help="Show version and exit",
+        ),
+        verbose: bool = typer.Option(False, "--verbose", "-v", help="More logs"),
+        dry_run: bool = typer.Option(False, "--dry-run", help="Print actions without changing files"),
+        local: bool = typer.Option(False, "--local", "-l", help="Use local proto.root instead of git source"),
+        config_path: Path = typer.Option(
+            Path("wizard/codegen.yaml"),
+            "--config", "-c",
+            exists=True,
+            file_okay=True,
+            dir_okay=False,
+            readable=True,
+            help="Path to codegen config YAML",
+        ),
+):
+    ctx.obj = Ctx(verbose=verbose, dry_run=dry_run, local=local, config_path=config_path)
+
+
+@app.command(help="Generate code from protos")
+def generate(ctx: typer.Context):
+    _print_verbose_enabled(ctx)
+    config = _load_config(ctx.obj.config_path, ctx)
+    proto_root = resolve_proto_root(config, use_local=ctx.obj.local)
+    files = discover_proto_files(proto_root, config)
+    if ctx.obj.verbose:
+        _print_files_table(files, proto_root)
+
+    fds_path, tmp_dir = build_descriptor_set(config, proto_root, files, ctx.obj.verbose)
+    print_fds_content(fds_path, ctx, console)
+    fds = load_fds(fds_path)
+    jinja_ctx = build_context(config, fds)
+    print_build_context(jinja_ctx, ctx, console)
+    plan = render_all(config, jinja_ctx)
+
+    apply_plan(plan, ctx.obj.dry_run, ctx.obj.verbose)
+
+    # Cleanup Tmp Dir
+    if tmp_dir:
+        shutil.rmtree(tmp_dir)
+
+    if ctx.obj.dry_run:
+        console.print("[bold green]✓[/] Dry run finished")
+    else:
+        console.print("[bold green]✓[/] Generated code")
+
+
+@app.command("list-protos", help="List available protos")
+def list_protos(ctx: typer.Context):
+    _print_verbose_enabled(ctx)
+    config = _load_config(ctx.obj.config_path, ctx)
+    proto_root = resolve_proto_root(config, use_local=ctx.obj.local)
+    files = discover_proto_files(proto_root, config)
+    _print_files_table(files, proto_root)
+
+
+@app.command(help="Validate the jinja2 templates, resolves proto root (git checkout), runs protoc descriptor build if needed, prints missing filters/variables")
+def validate(ctx: typer.Context):
+    _print_verbose_enabled(ctx)
+    config = _load_config(ctx.obj.config_path, ctx)
+    proto_root = resolve_proto_root(config, use_local=ctx.obj.local)
+    files = discover_proto_files(proto_root, config)
+
+    if ctx.obj.verbose:
+        _print_files_table(files, proto_root)
+
+    fds_path, tmp_dir = build_descriptor_set(config, proto_root, files, ctx.obj.verbose)
+    print_fds_content(fds_path, ctx, console)
+    fds = load_fds(fds_path)
+    jinja_ctx = build_context(config, fds)
+    print_build_context(jinja_ctx, ctx, console)
+    render_all(config, jinja_ctx)
+
+    # Cleanup Tmp Dir
+    if tmp_dir:
+        shutil.rmtree(tmp_dir)
+
+    console.print("[bold green]✓[/] Validation successful")
+
+
+def _print_files_table(files, proto_root):
+    table = Table(title="Available proto schemas")
+    table.add_column("Name", style="bold")
+    table.add_column("Path", style="dim")
+    for f in files:
+        table.add_row(f.stem, f.relative_to(proto_root).as_posix())
+    console.print(table)
+
+
+def _print_verbose_enabled(ctx: typer.Context):
+    if ctx.obj.verbose:
+        console.print("[dim]Verbose enabled[/]")
+
+def _load_config(path: Path, ctx: typer.Context) -> CodegenConfig:
+    config = load_config(path, console)
+    if ctx.obj.verbose:
+        console.print(
+            Panel(
+                Pretty(
+                    # pydantic v2
+                    config.model_dump(mode="python"),
+                    expand_all=True,
+                ),
+                title=f"[bold]Loaded config[/] [dim]{path}[/]",
+                border_style="cyan",
+            )
+        )
+    return config
+
+def main():
+    app()
+
+if __name__ == "__main__":
+    main()

core/__init__.py

@@ -0,0 +1,22 @@
+"""
+Core functionality for wizard code generation.
+
+Contains the main business logic for configuration, rendering, and writing files.
+"""
+
+from .config import CodegenConfig, load_config
+from .renderer import render_all, PlanItem
+from .writer import apply_plan
+from .context_builder import build_context, print_build_context
+from .filter import where_ok
+
+__all__ = [
+    "CodegenConfig",
+    "load_config",
+    "render_all",
+    "PlanItem",
+    "apply_plan",
+    "build_context",
+    "print_build_context",
+    "where_ok"
+]

core/config.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import typer
+import yaml
+from pathlib import Path
+from typing import Any
+from enum import Enum
+from pydantic import BaseModel, ConfigDict, Field, ValidationError
+from rich.console import Console
+from rich.panel import Panel
+
+class ForEachMode(str, Enum):
+    FILE = "file"
+    MESSAGE = "message"
+    ENUM = "enum"
+    SERVICE = "service"
+
+class WriteMode(str, Enum):
+    OVERWRITE = "overwrite"
+    APPEND = "append"
+    WRITE_ONCE = "write-once"
+
+class EqualsKV(BaseModel):
+    key: str
+    value: Any
+
+# A single predicate
+class Predicate(BaseModel):
+    # dotted selector like "name", "package", "file", "full_name"
+    # or "option" for descriptor options/custom options you extract into context
+    name: str | None = None
+    package: str | None = None
+    file: str | None = None
+    full_name: str | None = None
+    option_equals: EqualsKV | None = Field(default=None, alias="option.equals")
+
+    model_config = {"populate_by_name": True}
+
+class Where(BaseModel):
+    all: list[Predicate] = Field(default_factory=list)   # AND
+    any: list[Predicate] = Field(default_factory=list)   # OR
+    not_: list[Predicate] = Field(default_factory=list, alias="not")  # NOT
+
+    model_config = {"populate_by_name": True}
+
+class RenderTarget(BaseModel):
+    template: str
+    output: str
+    mode: WriteMode = WriteMode.OVERWRITE
+    for_each: ForEachMode | None = None
+    where: Where | None = None
+
+class ProtoSource(BaseModel):
+    git: str
+    ref: str = "latest-tag"  # Default to latest semver tag
+    fds: str | None = None
+    include_info: bool = True
+
+class ProtoConfig(BaseModel):
+    cache_dir: str
+    includes: list[str] = Field(default_factory=list)
+    files: list[str] = Field(default_factory=list)
+    root: str | None = ""
+    source: ProtoSource | None = None
+
+class TargetConfig(BaseModel):
+    templates: str
+    out: str
+    render: list[RenderTarget] = Field(default_factory=list)
+
+class HooksConfig(BaseModel):
+    root: str = "wizard"
+    module: str | None = None
+
+class CodegenConfig(BaseModel):
+    # nice to have: forbid typos in YAML keys
+    model_config = ConfigDict(extra="forbid")
+
+    proto: ProtoConfig
+    targets: dict[str, TargetConfig]
+    hooks: HooksConfig = Field(default_factory=HooksConfig)
+
+def _format_pydantic_errors(e: ValidationError) -> str:
+    lines: list[str] = []
+    for err in e.errors():
+        loc = ".".join(str(x) for x in err.get("loc", [])) or "<root>"
+        msg = err.get("msg", "Invalid value")
+        lines.append(f"• {loc}: {msg}")
+    return "\n".join(lines)
+
+def _die_config(message: str, console: Console, code: int = 2,) -> typer.Never:
+    console.print(Panel(message, title="[bold red]Config error[/]", border_style="red"))
+    raise typer.Exit(code=code)
+
+def load_config(path: Path, console: Console) -> CodegenConfig:
+    # 1) file read
+    try:
+        raw = path.read_text()
+    except FileNotFoundError:
+        _die_config(f"[bold]{path}[/] not found.\n\nPass a config with [cyan]--config[/].", console)
+    except OSError as e:
+        _die_config(f"Could not read [bold]{path}[/]: {e}", console)
+
+    # 2) YAML parse
+    try:
+        data: dict[str, Any] = yaml.safe_load(raw) or {}
+    except yaml.YAMLError as e:
+        _die_config(f"Invalid YAML in [bold]{path}[/]:\n\n[dim]{e}[/]", console)
+
+    # 3) Pydantic validate
+    try:
+        return CodegenConfig.model_validate(data)
+    except ValidationError as e:
+        problems = _format_pydantic_errors(e)
+        hint = (
+            "[bold]Expected top-level keys[/]\n"
+            "• proto\n"
+            "• targets\n"
+            "• hooks (optional)\n\n"
+            "[bold]Problems[/]\n"
+            f"{problems}"
+        )
+        _die_config(f"In [bold]{path}[/]:\n\n{hint}", console)

core/context_builder.py

@@ -0,0 +1,321 @@
+from __future__ import annotations
+from pathlib import Path
+from collections import defaultdict, deque
+from typing import Any, Dict, List
+
+import typer
+from google.protobuf import descriptor_pb2
+from rich.console import Console
+from rich.panel import Panel
+from rich.pretty import Pretty
+
+from .config import CodegenConfig
+from utils import Name
+
+
+def topo_order(files_by_name: dict[str, descriptor_pb2.FileDescriptorProto]) -> list[str]:
+    """Topologically sort proto files based on their dependencies."""
+    deps = {n: [d for d in f.dependency if d in files_by_name] for n, f in files_by_name.items()}
+    indeg = {n: 0 for n in files_by_name}
+    rev = defaultdict(list)
+
+    for n, ds in deps.items():
+        for d in ds:
+            indeg[n] += 1
+            rev[d].append(n)
+
+    q = deque([n for n, deg in indeg.items() if deg == 0])
+    out = []
+    while q:
+        n = q.popleft()
+        out.append(n)
+        for m in rev[n]:
+            indeg[m] -= 1
+            if indeg[m] == 0:
+                q.append(m)
+    return out
+
+
+def _pkg_path(pkg: str) -> str:
+    """Convert package name to path format."""
+    return pkg.replace(".", "/") if pkg else ""
+
+
+def _build_full_name(package: str, name: str) -> str:
+    """Build fully qualified name for a proto type."""
+    return f".{package}.{name}" if package else f".{name}"
+
+
+def _build_field_data(field: Any) -> Dict[str, Any]:
+    """Build field data dictionary from protobuf field descriptor."""
+    return {
+        "name": Name(field.name),
+        "number": field.number,
+        "label": int(field.label),
+        "type": int(field.type),
+        "type_name": field.type_name,
+        "json_name": field.json_name,
+        "field": field
+    }
+
+
+def _build_method_data(method: Any) -> Dict[str, Any]:
+    """Build method data dictionary from protobuf method descriptor."""
+    return {
+        "name": Name(method.name),
+        "input_type": method.input_type,
+        "output_type": method.output_type,
+        "client_streaming": bool(method.client_streaming),
+        "server_streaming": bool(method.server_streaming),
+        "options": method.options,
+    }
+
+
+def _build_enum_value_data(value: Any) -> Dict[str, Any]:
+    """Build enum value data dictionary from protobuf enum value descriptor."""
+    return {"name": value.name, "number": value.number}
+
+
+def _update_type_index(type_index: Dict[str, Any], full_name: str, kind: str, file_name: str, package: str) -> None:
+    """Update the type index with a new type entry."""
+    type_index[full_name] = {"kind": kind, "file": file_name, "package": package}
+
+
+def _process_nested_enum(
+        enum: Any,
+        parent_full_name: str,
+        file_name: str,
+        package: str,
+        type_index: Dict[str, Any],
+        enum_index: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Process a nested enum within a message."""
+    full_name = f"{parent_full_name}.{enum.name}"
+    enum_values = [_build_enum_value_data(value) for value in enum.value]
+
+    enum_data = {
+        "name": Name(enum.name),
+        "full_name": full_name,
+        "file": file_name,
+        "package": package,
+        "enum_values": enum_values,  # Use enum_values to avoid conflict with dict.values()
+    }
+
+    _update_type_index(type_index, full_name, "enum", file_name, package)
+    enum_index[full_name] = enum_data
+
+    return enum_data
+
+
+def _process_message_recursive(
+        message: Any,
+        parent_full_name: str,
+        file_name: str,
+        package: str,
+        type_index: Dict[str, Any],
+        message_index: Dict[str, Any],
+        enum_index: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Process a message and its nested types recursively."""
+    # Build full name for this message
+    if parent_full_name:
+        full_name = f"{parent_full_name}.{message.name}"
+    else:
+        full_name = _build_full_name(package, message.name)
+
+    fields = [_build_field_data(field) for field in message.field]
+
+    # Process nested messages recursively
+    nested_messages = []
+    for nested_msg in message.nested_type:
+        nested_data = _process_message_recursive(
+            nested_msg,
+            full_name,
+            file_name,
+            package,
+            type_index,
+            message_index,
+            enum_index
+        )
+        nested_messages.append(nested_data)
+
+    # Process nested enums
+    nested_enums = []
+    for nested_enum in message.enum_type:
+        nested_enum_data = _process_nested_enum(
+            nested_enum,
+            full_name,
+            file_name,
+            package,
+            type_index,
+            enum_index
+        )
+        nested_enums.append(nested_enum_data)
+
+    message_data = {
+        "name": Name(message.name),
+        "full_name": full_name,
+        "fields": fields,
+        "file": file_name,
+        "package": package,
+        "nested_messages": nested_messages,
+        "nested_enums": nested_enums,
+    }
+
+    _update_type_index(type_index, full_name, "message", file_name, package)
+    message_index[full_name] = message_data
+
+    return message_data
+
+
+def _process_messages(
+        proto_file: Any,
+        file_name: str,
+        package: str,
+        type_index: Dict[str, Any],
+        message_index: Dict[str, Any],
+        enum_index: Dict[str, Any]
+) -> List[Dict[str, Any]]:
+    """Process all messages in a proto file, including nested types."""
+    messages = []
+
+    for message in proto_file.message_type:
+        message_data = _process_message_recursive(
+            message,
+            "",  # No parent for top-level messages
+            file_name,
+            package,
+            type_index,
+            message_index,
+            enum_index
+        )
+        messages.append(message_data)
+
+    return messages
+
+
+def _process_enums(
+        proto_file: Any,
+        file_name: str,
+        package: str,
+        type_index: Dict[str, Any],
+        enum_index: Dict[str, Any]
+) -> List[Dict[str, Any]]:
+    """Process all enums in a proto file."""
+    enums = []
+
+    for enum in proto_file.enum_type:
+        full_name = _build_full_name(package, enum.name)
+        enum_values = [_build_enum_value_data(value) for value in enum.value]
+
+        enum_data = {
+            "name": Name(enum.name),
+            "full_name": full_name,
+            "file": file_name,
+            "package": package,
+            "enum_values": enum_values,  # Use enum_values to avoid conflict with dict.values()
+        }
+
+        enums.append(enum_data)
+        _update_type_index(type_index, full_name, "enum", file_name, package)
+        enum_index[full_name] = enum_data
+
+    return enums
+
+
+def _process_services(
+        proto_file: Any,
+        file_name: str,
+        package: str,
+        service_index: Dict[str, Any]
+) -> List[Dict[str, Any]]:
+    """Process all services in a proto file."""
+    services = []
+
+    for service in proto_file.service:
+        full_name = _build_full_name(package, service.name)
+        methods = [_build_method_data(method) for method in service.method]
+
+        service_data = {
+            "name": Name(service.name),
+            "methods": methods,
+            "file": file_name,
+            "package": package
+        }
+
+        services.append(service_data)
+        service_index[full_name] = service_data
+
+    return services
+
+
+def _process_proto_file(
+        proto_file: Any,
+        file_name: str,
+        type_index: Dict[str, Any],
+        message_index: Dict[str, Any],
+        enum_index: Dict[str, Any],
+        service_index: Dict[str, Any]
+) -> Dict[str, Any]:
+    """Process a single proto file and return its context data."""
+    package = proto_file.package
+    base_name = Path(file_name).stem
+
+    messages = _process_messages(proto_file, file_name, package, type_index, message_index, enum_index)
+    enums = _process_enums(proto_file, file_name, package, type_index, enum_index)
+    services = _process_services(proto_file, file_name, package, service_index)
+
+    return {
+        "name": Name(file_name),
+        "basename": base_name,
+        "package": package,
+        "package_path": _pkg_path(package),
+        "imports": list(proto_file.dependency),
+        "messages": messages,
+        "enums": enums,
+        "services": services,
+        "options": proto_file.options,
+    }
+
+
+def build_context(cfg: CodegenConfig, fds: descriptor_pb2.FileDescriptorSet) -> Dict[str, Any]:
+    """Build the complete context for code generation from proto file descriptors."""
+    files_by_name = {f.name: f for f in fds.file}
+    ordered_files = topo_order(files_by_name)
+
+    # Initialize indexes
+    type_index = {}
+    message_index = {}
+    enum_index = {}
+    service_index = {}
+    ctx_files = []
+
+    # Process each file in topological order
+    for file_name in ordered_files:
+        proto_file = files_by_name[file_name]
+        file_context = _process_proto_file(
+            proto_file, file_name, type_index, message_index, enum_index, service_index
+        )
+        ctx_files.append(file_context)
+
+    return {
+        "proto_root": cfg.proto.root,
+        "files": ctx_files,
+        "types": type_index,
+        "message": message_index,
+        "enum": enum_index,
+        "service": service_index,
+    }
+
+
+def print_build_context(ctx: Dict[str, Any], typer_ctx: typer.Context, console: Console,
+                        title: str = "Jinja2 context") -> None:
+    """Print the build context if verbose mode is enabled."""
+    if typer_ctx.obj.verbose:
+        console.print(
+            Panel(
+                Pretty(ctx, expand_all=True),
+                title=f"[bold]{title}[/]",
+                border_style="cyan",
+            )
+        )