bindings.py

"""Generate runtime handles and type stubs from static calculation definitions."""

from __future__ import annotations

import argparse
import json
import keyword
import os
import tempfile
from pathlib import Path

from .definitions import Definitions
from .source import Capture, SourceError, authoring_name


def generate(directory: Path, *, check: bool = False) -> dict:
    root = directory.resolve(strict=True)
    sources = sorted(root.rglob("*.cso.py"))
    if not sources:
        raise SourceError("MISSING_SOURCE", "No .cso.py calculations found")
    if not root.is_dir():
        raise SourceError("INVALID_DEPENDENCY_PATH", "Bindings require a directory")
    output = root / "_cso_bindings"
    if output.is_symlink() or any(path.is_symlink() for path in output.rglob("*")):
        raise SourceError(
            "INVALID_DEPENDENCY_PATH",
            "Generated bindings cannot contain symbolic links",
        )
    files: dict[Path, str] = {output / "__init__.py": "", output / "py.typed": ""}
    for source in sources:
        relative = source.relative_to(root)
        parts = [*relative.parts[:-1], relative.name[:-7]]
        if any(not part.isidentifier() or keyword.iskeyword(part) for part in parts):
            raise SourceError(
                "INVALID_BINDING_NAME",
                f"Binding module needs identifier path components: {relative}",
            )
        capture = Capture(source)
        capture.root = root
        capture.for_bindings = True
        module = capture.load(source)
        definitions = Definitions(capture)
        runtime = [
            "# Generated by cso bindings. Do not edit.\n",
            "from cso_python import load_calculation\n\n",
        ]
        stub = [
            "# Generated by cso bindings. Do not edit.\n",
            "from typing import TypedDict\n\n",
        ]
        target = output.joinpath(*parts).with_suffix(".py")
        parent = target.parent
        while parent != output:
            files[parent / "__init__.py"] = ""
            parent = parent.parent
        names = set(module.functions)
        reserved = {"TypedDict", "load_calculation"} | {
            f"_{name}_Outputs" for name in names
        }
        if names & reserved:
            raise SourceError(
                "BINDING_OUTPUT_CONFLICT",
                f"Calculation names collide with binding helpers: {sorted(names & reserved)}",
            )
        for name, fn in module.functions.items():
            if not any(
                authoring_name(d.func) == "calculation" for d in fn.decorator_list
            ):
                continue
            definition = definitions.get(source, name)
            path = os.path.relpath(source, target.parent)
            runtime.append(
                f"{name} = load_calculation({path!r}, function={name!r}, fingerprint={definition.fingerprint!r})\n"
            )
            output_type = f"_{name}_Outputs"
            fields = ", ".join(
                f"{key!r}: {spec.annotation.numeric_type}"
                for key, spec in definition.outputs.items()
            )
            stub.append(f"{output_type} = TypedDict({output_type!r}, {{{fields}}})\n")
            parameters = ", ".join(
                f"{key}: {spec.documented.numeric_type}"
                + (f" = {spec.default_value!r}" if spec.default is not None else "")
                for key, spec in definition.parameters.items()
            )
            stub.append(
                f"def {name}({('*, ' + parameters) if parameters else ''}) -> {output_type}: ...\n\n"
            )
        if target in files or target.with_suffix(".pyi") in files:
            raise SourceError(
                "BINDING_OUTPUT_CONFLICT", f"Colliding binding module {target}"
            )
        if target.name == "__init__.py":
            raise SourceError(
                "BINDING_OUTPUT_CONFLICT",
                "__init__.cso.py collides with a generated package",
            )
        files[target] = "".join(runtime)
        files[target.with_suffix(".pyi")] = "".join(stub)
    expected = set(files)
    existing = (
        {p for p in output.rglob("*") if p.is_file() and "__pycache__" not in p.parts}
        if output.exists()
        else set()
    )
    stale = [
        str(p.relative_to(root))
        for p in sorted(expected | existing)
        if p not in files or not p.exists() or p.read_text() != files[p]
    ]
    if not check:
        # Check ownership before replacing or removing an existing generated file.
        for path in existing:
            if not (
                path.name in ("__init__.py", "py.typed") and not path.read_bytes()
            ) and not path.read_text().startswith("# Generated by cso bindings."):
                raise SourceError(
                    "BINDING_OUTPUT_CONFLICT",
                    f"Refusing to overwrite authored file {path}",
                )
        for path, content in files.items():
            path.parent.mkdir(parents=True, exist_ok=True)
            descriptor, temporary = tempfile.mkstemp(
                dir=path.parent, prefix=".cso-bindings-"
            )
            try:
                with os.fdopen(descriptor, "w") as stream:
                    stream.write(content)
                os.replace(temporary, path)
            finally:
                if os.path.exists(temporary):
                    os.unlink(temporary)
        for path in existing - expected:
            path.unlink()
    return {
        "ok": not check or not stale,
        "command": "bindings",
        "check": check,
        "directory": str(root),
        "files": [str(p.relative_to(root)) for p in sorted(files)],
        "stale": stale,
    }


def bindings_from_argv(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(
        prog="python -m cso_python bindings", description=__doc__
    )
    parser.add_argument("directory", type=Path)
    parser.add_argument(
        "--check",
        action="store_true",
        help="Report stale bindings without writing files",
    )
    args = parser.parse_args(argv)
    try:
        result = generate(args.directory, check=args.check)
    except (SourceError, OSError, ValueError) as error:
        result = {
            "ok": False,
            "diagnostics": [
                error.diagnostic
                if isinstance(error, SourceError)
                else {"code": "BINDINGS_FAILED", "message": str(error)}
            ],
        }
    print(json.dumps(result, ensure_ascii=False, allow_nan=False))
    return 0 if result["ok"] else 1