configure

#!/usr/bin/env python

from __future__ import absolute_import, division, print_function

import subprocess
import os
from datetime import datetime
import argparse
import sys
from download_data import ftp_download

here = os.path.abspath(os.path.dirname(__file__))


def get_current_time():
    now = datetime.now().strftime("%H:%M:%S")
    return "({})".format(now)


def detect_platform():
    import platform

    p = {}
    p["system"] = platform.system().lower()
    p["node"] = platform.node().lower()
    return p


def detect_existing_executable(name):
    """ Check Whether `name` is on PATH. """
    from distutils.spawn import find_executable

    return find_executable(name)


def run_script(script, my_env, return_output=False):

    process_output = subprocess.Popen(
        script, env=my_env, stdout=subprocess.PIPE, stderr=subprocess.PIPE
    )

    out, err = process_output.communicate()
    if process_output.returncode != 0:
        raise RuntimeError(
            "{} script failed: {} {} {}".format(
                script, process_output.returncode, out.strip().decode(), err.strip().decode()
            )
        )

    if return_output:
        return out


def check_env_name_conflict(output, env_name):
    import json

    output = json.loads(output)["envs"]

    f = lambda x: os.path.basename(x)
    envs = set(map(f, output))

    return env_name in envs


def prepare_env(args):

    installers = {
        "linux": "https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh",
        "darwin": "https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-x86_64.sh",
    }

    my_env = os.environ.copy()
    my_env["BASE_ENV_YML"] = os.path.join(here, "environments/env-conda-base.yml")
    my_env["TUTORIAL_ENV_YML"] = os.path.join(here, "environments/env-tutorial.yml")
    default_shell = my_env.get("SHELL", None)
    supported_shells = {"bash", "fish", "powershell", "tcsh", "xonsh", "zsh"}
    if default_shell:
        default_shell = default_shell.split("/")[-1]
        if default_shell in supported_shells:
            my_env["INIT_SHELL"] = default_shell

    conda_path = detect_existing_executable("conda")
    if conda_path is not None:
        print("*** Found an existing Conda installation in: {}".format(conda_path))
        print("*** {} Step 1/6: Skipping Conda installation".format(get_current_time()))
        my_env["INSTALL_DIR"] = os.path.abspath(os.path.join(os.path.dirname(conda_path), ".."))

    else:
        print("*** {} Step 1/6: Installing Miniconda".format(get_current_time()))
        plat = detect_platform()
        if plat["system"] in installers:
            my_env["INSTALLER"] = installers[plat["system"]]
            my_env["INSTALL_DIR"] = args.prefix

            install_script = os.path.join(here, "conda/install_conda")
            run_script(install_script, my_env)

        else:
            raise RuntimeError("Unsupported Platform...")

    print("*** {} Step 2/6: Updating `base` environment".format(get_current_time()))
    run_script(os.path.join(here, "conda/update_base_env"), my_env)
    print(("*** {} Step 3/6: Checking for environment name conflict").format(get_current_time()))
    output = run_script(os.path.join(here, "conda/check_env_conflict"), my_env, return_output=True)
    env_exist = check_env_name_conflict(output, "python-tutorial")

    if env_exist and args.clobber:
        print(
            (
                "*** {} Step 4/6: Updating existing `python-tutorial` environment "
                "(this can take several minutes)"
            ).format(get_current_time())
        )
        run_script(os.path.join(here, "conda/update_tutorial_env"), my_env)

    elif env_exist and not args.clobber:
        print(
            "\n\nUnable to proceed to the next steps because of an existing",
            "`python-tutorial` environment",
            "with clobber set to {}.".format(args.clobber),
        )

        print("\nTo address this issue, please try one of the following options:")
        print("\n(1)Create clone environment, and then remove original one:")
        print("\n    conda create --name new_name --clone python-tutorial")
        print("\n    conda remove --name python-tutorial")
        print("\nOr")
        print("\n(2)you can run the configure script with clobber set to True:")
        print("\n    ./setup/configure --clobber")
        sys.exit()

    else:
        print(
            (
                "*** {} Step 4/6: Creating `python-tutorial` environment "
                "(this can take several minutes)"
            ).format(get_current_time())
        )
        run_script(os.path.join(here, "conda/update_tutorial_env"), my_env)

    print(
        "*** {} Step 5/6: Running post build script for `base` environment".format(
            get_current_time()
        )
    )
    run_script(os.path.join(here, "conda/post_build_base"), my_env)
    print(
        (
            "*** {} Step 6/6: Running post build script for `python-tutorial` environment "
            "(this can take several minutes)"
        ).format(get_current_time())
    )
    my_env["CARTOPY_ASSET_SCRIPT"] = os.path.join(here, "download_cartopy_assets.py")
    run_script(os.path.join(here, "conda/post_build_tutorial"), my_env)

    print("*** {} Setup completed successfully.".format(get_current_time()))
    print("==> For changes to take effect, close and re-open your current shell. <==")


def download_data():
    output_dir = os.path.abspath(os.path.join(os.path.dirname(here), "data"))
    local_data_locations = {
        "glade": "/glade/work/abanihi/aletheia-data/tutorial-data",
        "hobart": "/ftp/archive/aletheia-data/tutorial-data",
    }
    plat = detect_platform()

    def create_symlink(local_data_location):
        print(
            "Creating symlink to existing/local tutorial data directory: {}".format(
                local_data_location
            )
        )
        dst = output_dir + "/"
        src = os.path.join(local_data_location, "*")
        cmd = " ".join(["ln", "-sf", src, dst])
        print(cmd)
        proc = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        proc.wait()

    if plat["node"].startswith("casper") or plat["node"].startswith("cheyenne"):
        local_data_location = local_data_locations["glade"]
        create_symlink(local_data_location)

    elif plat["node"].startswith("hobart"):
        local_data_location = local_data_locations["hobart"]
        create_symlink(local_data_location)

    else:
        ftp_download(directory="archive/aletheia-data/tutorial-data", output_dir=output_dir)


if __name__ == "__main__":
    plat = detect_platform()
    if plat["node"].startswith("casper") or plat["node"].startswith("cheyenne"):
        default_install_dir = "/glade/work/{}/miniconda3".format(os.environ["USER"])
    else:
        default_install_dir = "{}/miniconda3".format(os.environ["HOME"])

    parser = argparse.ArgumentParser(
        description="Set up tutorial environment.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--clobber",
        "-c",
        help="Whether to clobber existing environment",
        default=False,
        action="store_true",
    )
    parser.add_argument(
        "--download",
        "-d",
        default=False,
        help="Download tutorial data without setting environment up",
        action="store_true",
    )
    parser.add_argument(
        "--prefix", "-p", default=default_install_dir, help="Miniconda3 install location"
    )

    args = parser.parse_args()
    if args.download:
        download_data()
    else:
        prepare_env(args)
        download_data()