make_env.py

"""A simple script to generate enviroment.yml files from pyproject.toml."""

import argparse
import fnmatch
import platform
import re
from pathlib import Path
from sys import version_info
from typing import List

if version_info.minor >= 11:
    from tomllib import load
else:
    from tomli import load

_FILE_DIR = Path(__file__).parent


# our quick and dirty implementation of recursive depedencies
def _parse_profile(profile_str: str, opt_deps: dict, project_name: str) -> List[str]:
    if profile_str is None or profile_str == "":
        return []

    pat = re.compile(r"\s*" + project_name + r"\[(.*)\]\s*")
    parsed = []
    queue = [f"{project_name}[{x.strip()}]" for x in profile_str.split(",")]
    while len(queue) > 0:
        dep = queue.pop(0)
        if dep == "":
            continue
        m = pat.match(dep)
        if m:
            # if we match the patern, all list elts have to be dependenciy groups
            dep_groups = [d.strip() for d in m.groups(0)[0].split(",")]
            unknown_dep_groups = set(dep_groups) - set(opt_deps.keys())
            if len(unknown_dep_groups) > 0:
                raise RuntimeError(f"unknown dependency group(s): {unknown_dep_groups}")
            queue.extend(dep_groups)
            continue

        if dep in opt_deps:
            queue.extend([x.strip() for x in opt_deps[dep]])
        else:
            parsed.append(dep)

    return parsed


parser = argparse.ArgumentParser()

parser.add_argument("profile", default="dev", nargs="?")
parser.add_argument("--output", "-o", default="environment.yml")
parser.add_argument("--channels", "-c", default=None)
parser.add_argument("--name", "-n", default=None)
parser.add_argument("--py-version", "-p", default=None)
args = parser.parse_args()

#
with open(Path(_FILE_DIR, "pyproject.toml"), "rb") as f:
    toml = load(f)
deps = toml["project"]["dependencies"]
opt_deps = toml["project"]["optional-dependencies"]
project_name = toml["project"]["name"]
# specific conda_install settings
install_config = toml["tool"].get("make_env", {})
conda_only = install_config.get("conda_only", [])
deps_not_in_conda = install_config.get("deps_not_in_conda", [])
channels = install_config.get("channels", ["conda-forge"])
if args.channels is not None:
    channels.extend(args.channels.split(","))
    channels = list(set(channels))

# parse environment name
name = args.name
if name is None:
    name = project_name.split("_")[1]
    if args.profile:
        name += f"_{args.profile}"
print(f"Environment name: {name}")

# parse dependencies groups and flavours
# "min" equals no optional dependencies
deps_to_install = deps.copy()
if args.profile not in ["", "min"]:
    extra_deps = _parse_profile(args.profile, opt_deps, project_name)
    deps_to_install.extend(extra_deps)

conda_deps = []
pip_deps = []
for dep in deps_to_install:
    if dep in deps_not_in_conda:
        pip_deps.append(dep)
    else:
        conda_deps.append(dep)
if args.py_version is not None:
    conda_deps.append(f"python=={args.py_version}")

pip_deps = sorted(list(set(pip_deps)))

# Make an exception for the build environment
if args.profile == "build":
    if platform.system().lower() == "windows":
        py = fnmatch.filter(conda_deps, "python*")
        gd = fnmatch.filter(conda_deps, "gdal*")
        np = fnmatch.filter(conda_deps, "numpy*")
        conda_deps.remove(*gd)
        conda_deps.remove(*np)
        if py:
            conda_deps.remove(*py)
        py = ["python==3.12.*"]
        pip_deps += conda_deps
        conda_deps = []
        if py:
            conda_deps += py
        pip_deps.append(
            "https://github.com/cgohlke/geospatial-wheels/releases/download/v2024.2.18/GDAL-3.8.4-cp312-cp312-win_amd64.whl",
        )
        pip_deps.append("numpy<2.0.0")

        for item in conda_only:
            im = fnmatch.filter(pip_deps, item)
            pip_deps.remove(*im)
            conda_deps.append(*im)

        pip_deps = sorted(list(set(pip_deps)))
        pip_deps.append("-e .")

# add pip as a conda dependency if we have pip deps
if len(pip_deps) > 0:
    conda_deps.append("pip")

# the list(set()) is to remove duplicates
conda_deps_to_install_string = "\n- ".join(sorted(list(set(conda_deps))))
channels_string = "\n- ".join(set(channels))

# create environment.yml
env_spec = f"""name: {name}

channels:
- {channels_string}

dependencies:
- {conda_deps_to_install_string}
"""
if len(pip_deps) > 0:
    pip_deps_to_install_string = "\n  - ".join(pip_deps)
    env_spec += f"""- pip:
  - {pip_deps_to_install_string}
"""

with open(Path(_FILE_DIR, args.output), "w") as out:
    out.write(env_spec)