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
import unicodedata
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",
        )
    modules: dict[Path, tuple[str, ...]] = {}
    packages: dict[tuple[str, ...], Path] = {}
    for source in sources:
        relative = source.relative_to(root)
        parts = (*relative.parts[:-1], relative.name[:-7])
        if parts[0] == "_external":
            raise SourceError(
                "BINDING_OUTPUT_CONFLICT",
                "_external is reserved for generated external handles",
            )
        if any(
            not part.isidentifier()
            or keyword.iskeyword(part)
            or unicodedata.normalize("NFKC", part) != part
            for part in parts
        ):
            raise SourceError(
                "INVALID_BINDING_NAME",
                f"Binding module needs identifier path components: {relative}",
            )
        modules[source] = parts
        for length in range(1, len(parts)):
            packages.setdefault(parts[:length], source)
    # Python prefers a package over a module with the same import name. Check
    # the complete plan before capturing definitions or touching existing files.
    for source, parts in modules.items():
        if parts in packages:
            child = packages[parts].relative_to(root)
            raise SourceError(
                "BINDING_OUTPUT_CONFLICT",
                f"Binding module {source.relative_to(root)} conflicts with "
                f"package {'/'.join(parts)}/ required by {child}",
            )
    files: dict[Path, bytes] = {output / "__init__.py": b"", output / "py.typed": b""}
    for source, parts in modules.items():
        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"] = b""
            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 = Path(os.path.relpath(source, target.parent)).as_posix()
            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.declared.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).encode("utf-8")
        files[target.with_suffix(".pyi")] = "".join(stub).encode("utf-8")
        for declaration_module in list(capture.modules.values()):
            if not declaration_module.externals:
                continue
            relative_module = declaration_module.path.relative_to(root).with_suffix("")
            if any(
                not part.isidentifier()
                or keyword.iskeyword(part)
                or unicodedata.normalize("NFKC", part) != part
                for part in relative_module.parts
            ):
                raise SourceError(
                    "INVALID_BINDING_NAME",
                    "External declaration modules need identifier path components",
                )
            external_target = output / "_external" / relative_module.with_suffix(".py")
            parent = external_target.parent
            while parent != output:
                files.setdefault(parent / "__init__.py", b"")
                parent = parent.parent
            external_runtime = [
                "# Generated by cso bindings. Do not edit.\nfrom cso_python import load_external\n\n"
            ]
            external_stub = [
                "# Generated by cso bindings. Do not edit.\nfrom typing import TypedDict\n\n"
            ]
            names = set(declaration_module.externals)
            if names & (
                {"TypedDict", "load_external"} | {f"_{name}_Outputs" for name in names}
            ):
                raise SourceError(
                    "BINDING_OUTPUT_CONFLICT",
                    "External names collide with generated binding helpers",
                )
            for name in declaration_module.externals:
                definition = definitions.get(declaration_module.path, name)
                relative_path = Path(
                    os.path.relpath(declaration_module.path, external_target.parent)
                ).as_posix()
                external_runtime.append(
                    f"{name} = load_external({relative_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()
                )
                external_stub.append(
                    f"{output_type} = TypedDict({output_type!r}, {{{fields}}})\n"
                )
                parameters = ", ".join(
                    f"{key}: {spec.declared.numeric_type}"
                    + (f" = {spec.default_value!r}" if spec.default is not None else "")
                    for key, spec in definition.parameters.items()
                )
                external_stub.append(
                    f"def {name}({('*, ' + parameters) if parameters else ''}) -> {output_type}: ...\n\n"
                )
            files[external_target] = "".join(external_runtime).encode("utf-8")
            files[external_target.with_suffix(".pyi")] = "".join(external_stub).encode(
                "utf-8"
            )
    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_bytes() != 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_bytes().startswith(b"# 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, "wb") 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