snek/cli/snek

#!/usr/bin/env python

import argparse
import json
import shutil
import subprocess
from collections import defaultdict
from datetime import datetime
from pathlib import Path


TS_FMT = "%Y%m%d-%H%M%S"


class Conda:

    def __init__(self, root=None, exe=None):
        self.exe = exe or shutil.which("mamba") or shutil.which("conda")
        if not self.exe:
            raise SystemExit("neither conda nor mamba found")

        root = root or self.get_root()
        self.root = Path(root)


    def fork(self, name):
        envs = self.get_envs()
        if name not in envs:
            envs = ", ".join(envs)
            msg = f'"{name}" is not a known environment.\nChoose from: {envs}'
            raise SystemExit(msg)

        basename = name
        if "@" in name:
            basename, original_timestamp = name.rsplit("@", 1)
            original_timestamp = datetime.strptime(original_timestamp, TS_FMT)
            print(f'extracted basename "{basename}" from {name} (original timestamp: {original_timestamp})')

        timestamp = datetime.now().strftime(TS_FMT)
        new = f"{basename}@{timestamp}"
        self.clone(name, new)


    def bless(self, name):
        if "@" not in name:
            msg = f'"{name}" is not a forked environment'
            raise SystemExit(msg)

        basename, timestamp = name.rsplit("@", 1)

        try:
            timestamp = datetime.strptime(timestamp, TS_FMT)
        except ValueError as e:
            raise SystemExit(e)

        self.clone(name, basename)


    def print_envs(self):
        envs = self.get_envs()
        for i in envs:
            print("-", i)


    def get_envs(self):
        cmd = [self.exe, "env", "list", "--json"]
        info = run_and_parse(cmd)
        envs = info["envs"]
        envs = sorted(envs)

        if self.root:
            envs_root = self.root / "envs"
            envs = [Path(i) for i in envs]
            envs = [i.relative_to(envs_root) for i in envs if i.is_relative_to(envs_root)]
            envs = [str(i) for i in envs]

        return envs


    def print_root(self):
        print(self.root)


    def get_root(self):
        info = self.get_info()
        return info["conda_prefix"]


    def print_exe(self):
        print(self.exe)


    def print_info(self):
        info = self.get_info()
        length = maxstrlen(info.keys())
        for k, v in sorted(info.items()):
            k = str(k) + ":"
            v = str(v)
            if len(v) >= 100:
                v = "\n" + v + "\n"
            print(k.ljust(length), v)


    def get_info(self):
        cmd = [self.exe, "info", "--json"]
        info = run_and_parse(cmd)
        return info


    def clone(self, original, new):
        # conda create --name theclone --clone theoriginal # uses hardlinks
        # conda create --name thecopy  --copy  theoriginal # copies files
        print("clone:", original, "=>", new)
        cmd = [self.exe, "create", "--name", new, "--clone", original]
        subprocess.run(cmd)


    def get_list(self, env):
        cmd = [self.exe, "list", "--json", "--name", env]
        pkgs = run_and_parse(cmd)
        res = defaultdict(list)
        for p in pkgs:
            for k in ("name", "version", "build_string", "build_number", "channel"):
                res[k].append(p[k])
        return dict(res)





def run_and_parse(cmd):
    completed = subprocess.run(cmd, capture_output=True)

    stderr = completed.stderr
    if stderr:
        msg = stderr.decode().strip()
        raise SystemExit(msg)

    stdout = completed.stdout
    if stdout:
        return json.loads(stdout)


def maxstrlen(seq):
    return max(len(str(i)) for i in seq)



def handle_clargs():
    parser = argparse.ArgumentParser(description="🐍")

    parser.add_argument("--root", help="conda root folder / prefix (read from the current env, if not given)")
    parser.add_argument("--exe", help="conda or mamba executable (read from the current env, if not given)")

    subparsers = parser.add_subparsers(title="commands", dest="command", required=True)

    parser_envs  = subparsers.add_parser("envs",  help="print existing envs")
    parser_root  = subparsers.add_parser("root",  help="print root")
    parser_exe   = subparsers.add_parser("exe",   help="print exe")
    parser_info  = subparsers.add_parser("info",  help="print info")

    parser_fork  = subparsers.add_parser("fork",  help="fork a conda env")
    parser_bless = subparsers.add_parser("bless", help="bless a conda env")

    msg = "name of the conda env"
    parser_fork.add_argument("name", help=msg)
    parser_bless.add_argument("name", help=msg)

    clargs = parser.parse_args()
    return clargs





if __name__ == "__main__":
    clargs = handle_clargs()
    conda = Conda(root=clargs.root, exe=clargs.exe)

    dispatch = {
        "envs": conda.print_envs,
        "root": conda.print_root,
        "exe": conda.print_exe,
        "info": conda.print_info,
        "fork": lambda: conda.fork(clargs.name),
        "bless": lambda: conda.bless(clargs.name)
    }

    dispatch[clargs.command]()