Build a package directory with mamba, but only expose the requested package to PATH

wrapped_mamba_environment.py

#!/usr/bin/env python3

"""
Build a wrapped mamba environment

Mamba is used to install a package, but then only the parts of the mamba
environment that come from that package are mirrored into a parent
directory structure. This allows you to use mamba to create and manage
the environment, without completely clobbering every other version of
every dependency that might already exist on the users PATH.
"""

import argparse
import itertools
import json
import re
import subprocess
import sys
from pathlib import Path

BOLD = "\033[1m"
NC = "\033[0m"

parser = argparse.ArgumentParser(
    description="Build a mamba environment exposing only specific packages"
)
parser.add_argument(
    "output_directory", type=Path, help="Target location folder to write to"
)
parser.add_argument(
    "packages",
    nargs="+",
    help="A package to expose, and any additional packages. The first package will be used for version.",
    metavar="package",
)
parser.add_argument(
    "--moduledir",
    "-m",
    type=Path,
    help="The module directory to write a versioned modulefile into",
)
parser.add_argument(
    "--versioned",
    action="store_true",
    help="Explicitly write to a version-named subfolder of the output directory",
)
parser.add_argument(
    "--channel",
    "-c",
    help="Channels to use. If specified, will override default environment.",
    action="append",
)
parser.add_argument("--debug", action="store_true", help=argparse.SUPPRESS)
args = parser.parse_args()

# The passed packages specifier might include a version constraint. Work
# out what the "primary" package name is


def _package_name(specifier: str) -> str:
    if not (match := re.match(r"[\w\.-]+", specifier)):
        sys.exit(f"Error: Could not package name from '{specifier}'")
    return match.group()


primary_package_name = _package_name(args.packages[0])

print(f"Primary package is: {BOLD}{primary_package_name}{NC}")

channels = []
if args.channel:
    channels = ["--override-channels"] + list(
        itertools.chain.from_iterable(("-c", c) for c in args.channel)
    )

if args.versioned:
    # We need to work out what version to install, before we install it

    try:
        print("Running mamba to determine install version in advance...")
        out = subprocess.run(
            [
                "mamba",
                "create",
                "-dp",
                "ENV_nonexistent",
                *channels,
                "--json",
                *args.packages,
            ],
            capture_output=True,
            check=True,
        )
    except subprocess.CalledProcessError:
        sys.exit("Error: Failed to run mamba to determine package version")
    # Load the mamba output
    data = json.loads(out.stdout)
    matching = [x for x in data["actions"]["LINK"] if x["name"] == primary_package_name]
    if not matching:
        sys.exit(
            f"Error: Could not find expected package {primary_package_name} in potential environment"
        )

    primary_package_version = matching[0]["version"]
    # Rewrite the package specifier so we still get this exact version
    args.packages[0] = f"{primary_package_name}={primary_package_version}"

    print(f"Primary package version: {BOLD}{primary_package_version}{NC}")

    # Append this version to the target directory
    args.output_directory /= primary_package_version

args.output_directory = args.output_directory.resolve()
conda_dir = args.output_directory.resolve() / "conda_env"

# Make sure this target location exists
args.output_directory.mkdir(parents=True, exist_ok=True)

print(f"Installing wrapped mamba environment to: {BOLD}{args.output_directory}{NC}")

try:
    print(f"Running {BOLD}mamba create{NC} ...")
    result = subprocess.run(
        [
            "mamba",
            "create",
            "-yp",
            str(conda_dir),
            "--json",
            *channels,
            *args.packages,
        ],
        check=True,
        capture_output=True,
    )
    print("done.")
except subprocess.CalledProcessError as e:
    print("Error in mamba:", e.stdout, e.stderr)
    sys.exit("Error: Could not create mamba environment")

packages_metadata = {x["name"]: x for x in json.loads(result.stdout)["actions"]["LINK"]}


def _read_package_metadata(conda_dir: Path, package_dist: str) -> dict:
    return json.loads((conda_dir / "conda-meta" / f"{package_dist}.json").read_bytes())


# Read out the metadata dictionary for every package specified on run
metadata = {}
for package in [_package_name(x) for x in args.packages]:
    metadata[package] = _read_package_metadata(
        conda_dir, packages_metadata[package]["dist_name"]
    )

targets: dict[str, list[Path]] = {
    "bin": [],
    "man": [],
}

# Work out, from this, what paths we want to mirror
for path in itertools.chain.from_iterable(
    x["paths_data"]["paths"] for x in metadata.values()
):
    if path["_path"].startswith("bin/") and path["_path"].count("/") == 1:
        targets["bin"].append(Path(path["_path"]))
    elif path["_path"].startswith("man/"):
        targets["man"].append(Path(path["_path"]))
    elif path["_path"].startswith("share/man/"):
        targets["man"].append(Path(path["_path"]))

if targets["bin"]:
    print("Processing Binary targets:")
    bindir = args.output_directory / "bin"
    bindir.mkdir(exist_ok=True)
    for bin in targets["bin"]:
        dest_path = bindir / bin.name
        target_path = conda_dir / bin
        if dest_path.exists():
            print(f"    {dest_path} already exists")
        else:
            print(f"    {dest_path} → {target_path}")
            dest_path.symlink_to(target_path)

if targets["man"]:
    print("Processing manpage targets:")
    mandir = args.output_directory / "man"
    mandir.mkdir(exist_ok=True)
    for man in targets["man"]:
        if man.parts[:2] == ("share", "man"):
            dest = Path(*man.parts[2:])
        else:
            dest = Path(*man.parts[1:])
        # Make sure the target directory exists
        (mandir / dest.parent).mkdir(exist_ok=True, parents=True)
        dest_path = mandir / dest
        target_path = conda_dir / man
        if dest_path.exists():
            print(f"    {dest_path} already exists")
        else:
            print(f"    {dest_path} → {target_path}")
            dest_path.symlink_to(target_path)

# Build the Modulefile, so we can write it out
package_version = metadata[primary_package_name]["version"]

modulefile_parts = [
    f"""
#%Module
module-whatis "{primary_package_name} {package_version}"
""".lstrip()
]

if targets["bin"]:
    modulefile_parts.append(f"prepend-path PATH {args.output_directory}/bin")
if targets["man"]:
    modulefile_parts.append(f"prepend-path MANPATH {args.output_directory}/man")

if args.moduledir:
    modulefile = args.moduledir / package_version
    print(f"Writing ModuleFile {modulefile}")
    args.moduledir.mkdir(exist_ok=True, parents=True)
    modulefile.write_text("\n".join(modulefile_parts))

print("\nDone.")