myst-venv-pep723.py

#!/usr/bin/env python3
import re
import sys
import tomllib
import json
import argparse
import pathlib
import shutil
import subprocess


REGEX = r"(?m)^# /// (?P<type>[a-zA-Z0-9-]+)$\s(?P<content>(^#(| .*)$\s)+)^# ///$"


def parse_script_metadata(script: str) -> dict | None:
    name = "script"
    matches = list(
        filter(lambda m: m.group("type") == name, re.finditer(REGEX, script))
    )
    if len(matches) > 1:
        raise ValueError(f"Multiple {name} blocks found")
    elif len(matches) == 1:
        content = "".join(
            line[2:] if line.startswith("# ") else line[1:]
            for line in matches[0].group("content").splitlines(keepends=True)
        )
        return tomllib.loads(content)
    else:
        return None


def iter_nodes(mdast, type_):
    if mdast["type"] == type_:
        yield mdast

    if "children" in mdast:
        for child in mdast["children"]:
            yield from iter_nodes(child, type_)


def parse_document_spec(content):
    for node in iter_nodes(content["mdast"], "code"):
        metadata = parse_script_metadata(node["value"])
        if metadata is not None:
            break
    else:
        return None

    script_dependencies = metadata.get("dependencies", [])
    return [*script_dependencies, "ipykernel"]


VENV_PATH = pathlib.Path(".myst-venvs").absolute()


def main(argv=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("build_path", type=pathlib.Path)
    parser.add_argument(
        "-r", "--root-venv-path", type=pathlib.Path, default=pathlib.Path(sys.prefix)
    )
    args = parser.parse_args(argv)

    envs_path = VENV_PATH
    if envs_path.exists():
        shutil.rmtree(envs_path)

    root_venv_path = args.root_venv_path.absolute()

    for p in args.build_path.absolute().glob("**/*.json"):
        with open(p, "r") as f:
            content = json.load(f)

        kernel_name = content["frontmatter"].get("kernelspec", {}).get("name")
        if kernel_name is None:
            doc_name = content["location"].removeprefix("/")
            print(f"Missing kernel name in {doc_name}")
            continue

        spec = parse_document_spec(content)
        if spec is None:
            continue

        env_path = envs_path / kernel_name
        env_path.mkdir(exist_ok=True, parents=True)

        print(f"Provisioning environment for {p.name} at {env_path}")
        build_environment(kernel_name, spec, env_path, root_venv_path)


def build_environment(kernel_name, spec, env_path, root_prefix):
    subprocess.run([sys.executable, "-m", "venv", env_path], check=True)
    interp_path = env_path / "bin" / "python"

    subprocess.run([interp_path, "-m", "pip", "install", *spec, "-qqq"], check=True)

    subprocess.run(
        [
            interp_path,
            "-m",
            "ipykernel",
            "install",
            "--name",
            kernel_name,
            "--prefix",
            root_prefix,
        ],
        check=True,
    )


if __name__ == "__main__":
    main()