pycli

#!/usr/bin/env python3

import argparse
import functools
import os
import pathlib
import shutil
import subprocess
import sys
import types
import urllib.request


class _Registry:
    def __init__(self):
        self.registered = {}

    def __call__(self, function=None, name=None):
        if function is None and name is None:
            raise TypeError("Pass a function or name.")

        if function is None:
            return functools.partial(self, name=name)

        self.registered[function.__name__.replace("_", "-")] = function
        return function


_register = _Registry()


def build_requirements_name(layer, extension):
    if layer is None:
        return "requirements." + extension

    return layer + "-requirements." + extension


@_register
def clean(cfg):
    """Remove extraneous files."""
    paths = (
        [str(cfg.venv_path), ".coverage"]
        + list(pathlib.Path().glob(".coverage.*"))
        + ["dist"]
    )

    for path in paths:
        try:
            shutil.rmtree(path)
        except FileNotFoundError:
            pass


@_register
def init(cfg):
    """Set up a virtualenv, install requirements.txt, dev-requirements.txt, and current dir."""

    subprocess.run(
        ["virtualenv", "--python", sys.executable, str(cfg.venv_path)], check=True
    )
    if not pathlib.Path(build_requirements_name(None, "txt")).exists():
        raise FileNotFoundError("Run `lock` first, to create requirements.txt.")
    if pathlib.Path(build_requirements_name("dev", "txt")).exists():
        subprocess.run(
            [
                cfg.venv_path / "bin/pip",
                "install",
                "--requirement",
                build_requirements_name("dev", "txt"),
            ],
            check=True,
        )
    subprocess.run(
        [cfg.venv_path / "bin/pip", "install", "--requirement", build_requirements_name(None, "txt")],
        check=True,
    )
    subprocess.run(
        [cfg.venv_path / "bin/pip", "install", "--editable", "."], check=True
    )


@_register
def lock(cfg):
    """Use pip-compile to generate package hashes from setup.py and write them into requirements.txt."""
    subprocess.run([cfg.venv_path / "bin/pip", "install", "pip-tools"], check=True)
    combined = []
    for layer in [None, 'test', 'dev']:
        combined.append(layer)

        subprocess.run(
            [
                cfg.venv_path / "bin/pip-compile",
                "--generate-hashes",
                "--output-file",
                build_requirements_name(layer, 'txt'),
                *(
                    build_requirements_name(prefix, 'in')
                    for prefix in combined
                )
            ],
            check=True,
            env={
                **os.environ,
                "CUSTOM_COMPILE_COMMAND": "python {} lock".format(
                    pathlib.Path(__file__).name
                ),
            },
        )


@_register
def build(cfg):
    """Build source and binary distributions."""
    subprocess.run(
        [cfg.venv_path / "bin/python", "setup.py", "sdist", "bdist_wheel"], check=True
    )


@_register
def upload(cfg):
    """Upload the distributions to PyPI."""
    subprocess.run([cfg.venv_path / "bin/python", "-m", "pip", "install", "twine"])
    dists = [str(path) for path in pathlib.Path("dist").iterdir()]
    subprocess.run([cfg.venv_path / "bin/twine", "upload", *dists], check=True)


@_register
def bundle(cfg):
    """Bundle the package into a standalone unix executable."""
    lock(cfg)
    with open(build_requirements_name(None, "txt")) as f:
        requirements = [line.split()[0] for line in f if line[0].isalpha()]

    subprocess.run(
        [
            cfg.venv_path / "bin/pex",
            ".",
            *requirements,
            "-m",
            "desert",
            "-o" "desert.pex",
            "--disable-cache",
        ]
    )


def _get_default_venv_path():
    """Get the default path of the venv."""

    if (pathlib.Path().resolve() / "venv").exists():
        return pathlib.Path().resolve() / "venv"

    venv_path = os.environ.get("VENV_PATH")
    if venv_path:
        return pathlib.Path(venv_path)

    workon_home = os.environ.get("WORKON_HOME")
    if workon_home is not None:
        project_name = pathlib.Path(os.getcwd()).name
        return pathlib.Path(workon_home) / project_name

    return pathlib.Path().resolve() / "venv"


def cli():
    parser = argparse.ArgumentParser()

    parser.add_argument(
        "--venv",
        default=_get_default_venv_path(),
        type=pathlib.Path,
        help="Path of the venv directory. Defaults, in order: venv (if exists already), $VENV_PATH, $WORKON_HOME/[current directory name], venv",
    )

    subparsers = parser.add_subparsers(dest="command_name")
    for name, function in _register.registered.items():
        subparsers.add_parser(name, help=function.__doc__)

    args = parser.parse_args()
    if args.command_name is None:
        parser.print_help()
        sys.exit(2)

    function = _register.registered[args.command_name]
    cfg = types.SimpleNamespace()
    cfg.venv_path = args.venv
    function(cfg)


if __name__ == "__main__":
    cli()