setup_py_env.py

"""Setup python env in our build."""
import os
import re
import sys
import typing
from pathlib import Path

import automata
import print_site

IS_DEBUG = False


def is_nt() -> bool:
    """Check if NT."""
    return os.name == "nt"


def bindir() -> str:
    """Report scripts or bin."""
    return "Scripts" if is_nt() else "bin"


def ext() -> str:
    """Report python extension."""
    return ".exe" if is_nt() else ""


def fix_activate_path(VRdir: Path) -> None:
    """
    Fix activate path.

    :param VRdir:
    """
    new_lines = []
    fn = Path(VRdir) / bindir() / "activate"
    with open(fn) as fh:
        sexp = re.compile(r"VIRTUAL_ENV='(.+)'")
        for line in fh:
            # VIRTUAL_ENV='C:\Users\appsmith\asv\pyenv\versions\glue-run-dbg'
            m = sexp.search(line)
            if m:
                # convert slashes
                the_path = str(Path(m.group(1)).as_posix())
                # fix drive letter
                fix_path = f"/{the_path[0].lower()}{the_path[2:]}"
                line = f"VIRTUAL_ENV='{fix_path}'\n"
            new_lines.append(line)

    with open(fn, "w") as fh:
        print("".join(new_lines), file=fh)


def pyexe(PR: Path) -> Path:
    """Path of python exe."""
    return PR / bindir() / f"python{ext()}"


def ensure_pip(PR: Path) -> None:
    """Ensure pip."""
    pe = pyexe(PR)
    automata.sp_string(
        [pe, "-s", "-m", "ensurepip", "--default-pip", "--upgrade", "--verbose"]
    )


def upgrade_pip(PR: Path) -> None:
    """Upgrade pip."""
    pe = pyexe(PR)
    automata.sp_run(
        [pe, "-s", "-m", "pip", "install", "--upgrade", "--verbose pip"]
    )


def compose_venv_root(PR: Path) -> Path:
    """Compose a virtualenv."""
    # where do we put the venv?
    drive = "i:/" if is_nt() else "/i"
    vext = "dbg" if "debug" in str(PR) else "rel"
    pyenv_dir = Path(drive, "pyenv", "versions", f"glue-run-{vext}")
    print("PYENV_DIR={pyenv_dir}")
    return pyenv_dir


def make_venv(PR: Path, VRdir: Path) -> None:
    """Make a virtualenv."""
    venvexe = PR / bindir() / f"virtualenv{ext()}"
    pyexe = PR / bindir() / f"python{ext()}"
    automata.sp_string(
        [
            venvexe,
            VRdir,
            f"--python={pyexe}",
            "--verbose",
            "--always-copy",
            "--clear",
        ]
    )

    if is_nt():
        fix_activate_path(VRdir)


def activate_venv(VR: Path) -> None:
    """Activate a virtualenv."""
    # activate the venv - sets a few vars
    # https://stackoverflow.com/questions/436198/what-is-an-alternative-to-execfile-in-python-3
    activate_script = VR / bindir() / "activate_this.py"
    myglobals = dict(__file__=activate_script, __name__="__main__")

    print(f"activating virtualenv={VR} with {activate_script}")
    with open(activate_script, "rb") as f:
        code = compile(f.read(), activate_script, "exec")
        exec(code, myglobals)

    print_site.print_site()


def install_pkgs(
    pkglist: typing.List[str], PR: Path, do_upgrade: bool = True
) -> None:
    """Install packages."""
    pe = pyexe(PR)
    uparg = "" if not do_upgrade else "--upgrade"
    automata.sp_run(
        [pe.as_posix(), "-s", "-m", "pip", "install", uparg, "--verbose"]
        + pkglist
    )


def install_ports(portlist: typing.List[str], PR: Path = None) -> None:
    """Install ports."""
    drive = "i:/" if is_nt() else "/i"
    portboy = Path(drive, "ports", "scripts", "portboy.py")
    pe = pyexe(PR)
    automata.sp_run([pe.as_posix(), portboy.as_posix()] + portlist)


def install_virtualenv(PR: Path = None) -> None:
    """Install our virtualenv."""
    pkglist = [
        "virtualenv",
    ]
    install_pkgs(pkglist, PR, do_upgrade=False)


def install_our_pkgs(PR: Path = None) -> None:
    """Install our packages."""
    pkglist = [
        "docopt",
        "msgpack",
        "mashumaro",
        "mashuhelpa",
        "rpyc",
        "pyyaml",
        "sqlalchemy",
        "fdb",
        "graphql-core",
        "pyrsistent",
        "datetime",
        "snakemake",
        "twine",
        "region_profiler",
        "pypreprocessor",
    ]

    # portlist = ()

    install_pkgs(pkglist, PR)
    # install_ports(a, portlist, PR)


def copy_python_exe(PR: Path) -> None:
    """Copy python exe."""
    pydir = PR / bindir()
    pyexe = str(pydir / f"python{ext()}")
    lncmd = ["ln", "-s", pyexe]

    py3exe = pydir / f"python3{ext()}"
    if not py3exe.exists():
        # link python3 to python
        automata.sp_run([lncmd, py3exe])
    if is_nt():
        # link pythonw to python
        pywexe = pydir / f"pythonw{ext()}"
        if not pywexe.exists():
            automata.sp_run([lncmd, pywexe])


def fix_dll_search_path() -> None:
    """Fixup dll seatch path."""
    import win_fix_dlls

    win_fix_dlls.add_path_to_dll_search(True)


def fix_dll(fname: str, use_dbg_stem: str = "_d") -> str:
    """Fixup dll."""
    global IS_DEBUG
    suffix = ".dll" if is_nt() else ".so"
    prefix = "" if is_nt() else "lib"
    dbg_stem = use_dbg_stem if IS_DEBUG else ""
    return f"{prefix}{fname}{dbg_stem}{suffix}"


def fix_dll_list(
    base_list: typing.List[Path], use_dbg_stem: str = "_d"
) -> typing.List[Path]:
    """Fixup dll list."""
    return [p.parent / fix_dll(p.name, use_dbg_stem) for p in base_list]


def copy_ext_dlls(PR: Path) -> None:
    """Copy extenstion dlls."""
    if is_nt():
        # 3.8+ extensions and c-types don't search PATH!!!
        # so dlls have to be alongside the pyd
        dll_subdir = "bin" if is_nt() else "lib"
        plat_dir = Path(os.environ["ASV_PLAT_PORTS"])
        plat_dll_dir = plat_dir / dll_subdir
        d_debug_list = [
            plat_dll_dir / "libssl32MD",
            plat_dll_dir / "libcrypto32MD",
        ]
        debug_list = [
            plat_dll_dir / "zlib",
            plat_dll_dir / "libffi",
            plat_dll_dir / "libexpat",
            plat_dll_dir / "sqlite3-shared",
        ]
        dll_list = fix_dll_list(d_debug_list, "d") + fix_dll_list(debug_list)
        if IS_DEBUG:
            # include the release dlls!!!! something is goofed...
            dll_list.extend(fix_dll_list(debug_list, ""))

        automata.real_cp(dll_list, PR / "DLLs")


def do_setup(PR: Path) -> None:
    """Do the setup."""
    # fix_dll_search_path()

    print(f"setup_py_env::do_setup {PR}")

    # on Windows, pip and such will fish in the registry
    # path still has various python versions first
    # better all be at same level!
    print_site.print_site()

    # install into our build
    copy_python_exe(PR)
    copy_ext_dlls(PR)

    ensure_pip(PR)
    upgrade_pip(PR)
    install_virtualenv(PR)

    # create a virtualenv using python install
    # VR = compose_venv_root(PR)
    # make_venv(PR, VR)
    # activate_venv(VR)
    # upgrade_pip(VR)
    # install_our_pkgs(VR)


def usage() -> str:
    """Return usage string."""
    msg = f"""\
Usage: python {__file__} python_root
"""
    print(msg)
    sys.exit(-1)


def main(argv: typing.List[str] = None) -> int:
    """Run the main program."""
    global IS_DEBUG
    if argv is None:
        argv = sys.argv

    if len(argv) < 2:
        usage()
        return -1

    # caller to pass pyroot
    PR = Path(argv[1])
    IS_DEBUG = "debug" in str(PR)

    do_setup(PR)

    return 0


if __name__ == "__main__":
    sys.exit(main())