__init__.py

'''class Tool, class InstallMethod, and related subclasses and methods'''

__author__ = "dpark@broadinstitute.org,irwin@broadinstitute.org"

import collections
import json
import operator
import os
import re
import logging
import tempfile
import shutil
import shlex
import subprocess
import util.file
import util.misc

try:
    # Python 3.x
    from urllib.request import urlretrieve    # pylint: disable=E0611
    from urllib.parse import urlparse    # pylint: disable=E0611
except ImportError:
    # Python 2.x
    from urllib import urlretrieve # pylint: disable=E0611
    from urlparse import urlparse # pylint: disable=import-error

# Put all tool files in __all__
# allows "from tools import *" to import all tooles for testtools
__all__ = sorted(
    [
        filename[:-3]    # Remove .py
        for filename in os.listdir(os.path.dirname(__file__))    # tools directory
        if filename.endswith(
            '.py') and filename != '__init__.py' and filename not in [    # Add any files to exclude here:
    # e.g. 'sometool.py',
            ]
    ]
)
installed_tools = {}

_log = logging.getLogger(__name__)


def iter_leaf_subclasses(aClass):
    "Iterate over subclasses at all levels that don't themselves have a subclass"
    isLeaf = True
    for subclass in sorted(aClass.__subclasses__(), key=operator.attrgetter("__name__")):
        isLeaf = False
        for leafClass in iter_leaf_subclasses(subclass):
            if not getattr(leafClass, '_skiptest', False):
                yield leafClass
    if isLeaf:
        yield aClass

def all_tool_classes():
    return iter_leaf_subclasses(Tool)


def get_tool_by_name(name):
    if name not in installed_tools:
        raise NotImplementedError
    return installed_tools[name]


def skip_install_test(condition=None):
    '''Decorate the Tool class to skip the installation test.'''

    def decorator(klass):
        if callable(condition) and not condition():
            return klass
        klass._skiptest = True
        return klass

    return decorator


def is_osx():
    return os.uname()[0] == 'Darwin'


class Tool(object):
    ''' Base tool class that includes install machinery.

        TO DO: add something about dependencies..
    '''

    def __init__(self, install_methods=None):
        install_methods = install_methods or []

        self.install_methods = install_methods
        self.installed_method = None
        self.exec_path = None

    def is_installed(self):
        return (self.installed_method is not None)

    def install(self):
        if not self.is_installed():
            for m in self.install_methods:
                if not m.is_attempted():
                    m.attempt_install()
                if m.is_installed():
                    self.installed_method = m
                    self.exec_path = m.executable_path()
                    installed_tools[self.__class__.__name__] = self
                    break

    def get_install_methods(self):
        return self.install_methods

    def set_install_methods(self, methods):
        self.install_methods = methods

    def version(self):
        return None

    def executable_path(self):
        return self.exec_path

    def execute(self, *args):
        assert not os.system(self.exec_path + ' ' + args)

    def install_and_get_path(self):
        self.install()
        if self.executable_path() == None:
            raise NameError("unsuccessful in installing " + type(self).__name__)
        return self.executable_path()


class InstallMethod(object):
    ''' Base class for installation methods for a given tool.
        None of these methods should ever fail/error. attempt_install should
        return silently regardless of the outcome (is_installed must be
        called to verify success or failure).
    '''

    def __init__(self):
        self.attempts = 0

    def is_attempted(self):
        return self.attempts

    def attempt_install(self):    # Override _attempt_install, not this.
        self.attempts += 1
        self._attempt_install()

    def _attempt_install(self):
        raise NotImplementedError

    def is_installed(self):
        raise NotImplementedError

    def executable_path(self):
        raise NotImplementedError


class PrexistingUnixCommand(InstallMethod):
    ''' This is an install method that tries to find whether an executable
        binary already exists for free on the unix file system--it doesn't
        actually try to install anything.
    '''

    def __init__(self, path, verifycmd=None, verifycode=0, require_executability=True):
        self.path = path
        self.verifycmd = verifycmd
        self.verifycode = verifycode
        self.installed = False
        self.require_executability = require_executability
        InstallMethod.__init__(self)

    def _attempt_install(self):
        if os.access(self.path, (os.X_OK | os.R_OK) if self.require_executability else os.R_OK):
            if self.verifycmd:
                self.installed = (os.system(self.verifycmd) == self.verifycode)
            else:
                self.installed = True
        else:
            self.installed = False

    def is_installed(self):
        if not self.is_attempted():
            self.attempt_install()
        return self.installed

    def executable_path(self):
        return self.installed and self.path or None


class CondaPackageVersion(object):

    def __init__(self, version, build_type=None):
        self.version = version
        self.build_type = build_type

    def __cmp__(self, other):
        return self.__dict__ == other.__dict__

    def __repr__(self):
        if self.build_type:
            return '{}-{}'.format(self.version, self.build_type)
        else:
            return self.version

    def satisfies(self, spec):
        if spec.build_type:
            if self.build_type != spec.build_type:
                return False
        return self.version == spec.version

    @property
    def version_spec(self):
        if self.build_type:
            return '{}={}'.format(self.version, self.build_type)
        else:
            return self.version


class CondaPackage(InstallMethod):
    ''' This is an install method for tools that can be installed via
        conda.
    '''

    def __init__(
        self,
        package,
        channel="bioconda",
        executable=None,
        version="",
        verifycmd=None,
        verifycode=0,
        require_executability=True,
        env=None,
        env_root_path=None,
        conda_cache_path=None,
        patches=None,
        post_install_command=None,
        post_install_ret=0,
        post_verify_command=None,
        post_verify_ret=0
    ):
        # if the executable name is specifed, use it; otherwise use the package name
        self.executable = executable or package
        self.package = package
        self.channel = channel
        if type(version) == CondaPackageVersion:
            self.version = version
        else:
            self.version = CondaPackageVersion(version)

        self.post_install_command = post_install_command
        self.post_install_ret = post_install_ret

        # call the post-verification command.
        # Useful for copying in license files, building databases, etc.
        # The post-verify command is executed relative to the conda environment bin/
        # And has the active conda environment on the PATH
        self.post_verify_command = post_verify_command
        self.post_verify_ret = post_verify_ret
        self.post_verify_cmd_executed = False

        self.verifycmd = verifycmd
        self.verifycode = verifycode
        self.require_executability = require_executability
        self.patches = patches or []

        # if we have specified a conda env/root, use it
        # alternatively, use cms tools dir as default env root if os.environ["CONDA_PREFIX"] is not defined
        #
        # as it is always the path
        # CONDA_PREFIX is always a full path, but may not be present in older versions
        # CONDA_ENV_PATH may be used instead
        # in even older versions of conda, CONDA_DEFAULT_ENV can be used as another fallback
        self.env_path = None
        if "CONDA_PREFIX" in os.environ and len(os.environ["CONDA_PREFIX"]):
            #_log.debug('CONDA_PREFIX found')
            last_path_component = os.path.basename(os.path.normpath(os.environ["CONDA_PREFIX"]))
            self.env_path = os.path.dirname(os.environ["CONDA_PREFIX"]) if last_path_component == "bin" else os.environ["CONDA_PREFIX"]
        elif "CONDA_ENV_PATH" in os.environ and len(os.environ["CONDA_ENV_PATH"]):
            #_log.debug('CONDA_ENV_PATH found')
            last_path_component = os.path.basename(os.path.normpath(os.environ["CONDA_ENV_PATH"]))
            self.env_path = os.path.dirname(os.environ["CONDA_ENV_PATH"]) if last_path_component == "bin" else os.environ["CONDA_ENV_PATH"]
        elif "CONDA_DEFAULT_ENV" in os.environ and len(os.environ["CONDA_DEFAULT_ENV"]):
            #_log.debug('CONDA_PREFIX not found, using CONDA_DEFAULT_ENV')
            conda_env_path = os.environ.get('CONDA_DEFAULT_ENV')  # path to current conda environment
            if conda_env_path:
                if os.path.isdir(conda_env_path):
                    #_log.debug('Conda env found is specified as dir: %s' % conda_env_path)
                    conda_env_path = os.path.abspath(conda_env_path)
                    last_path_component = os.path.basename(os.path.normpath(conda_env_path))
                    self.env_path = os.path.dirname(last_path_component) if last_path_component == "bin" else conda_env_path
                else: # if conda env is an environment name, infer the path
                    #_log.debug('Conda env found is specified by name: %s' % conda_env_path)
                    result = util.misc.run_and_print(["conda", "env", "list", "--json"], silent=True, env=os.environ)
                    if result.returncode == 0:
                        command_output = result.stdout.decode("UTF-8")
                        data = json.loads(self._string_from_start_of_json(command_output))
                        if len(data) > 0:
                            if "envs" in data and len(data["envs"]):
                                for item in data["envs"]:
                                    if os.path.basename(os.path.realpath(item)) == conda_env_path:
                                        self.env_path = os.path.realpath(item)
                                        break

        # if the env is being overridden, or if we could not find an active conda env
        if env_root_path or env or not self.env_path:
            env_root_path = env_root_path or os.path.join(util.file.get_project_path(), 'tools', 'conda-tools')
            env = env or 'default'
            self.env_path = os.path.realpath(os.path.expanduser(
                os.path.join(env_root_path, env)))

        # set an env variable to the conda cache path. this env gets passed to the
        # the subprocess, and the variable instructs conda where to place its cache files
        conda_cache_path = conda_cache_path or os.path.join(util.file.get_project_path(), 'tools', 'conda-cache')
        self.conda_cache_path = os.path.realpath(os.path.expanduser(conda_cache_path))
        self.conda_env = os.environ
        old_envs_path = os.environ.get('CONDA_DEFAULT_ENV')
        self.conda_env["CONDA_ENVS_PATH"] = conda_cache_path+":"+os.path.dirname(self.env_path)

        #_log.info("Tool install conda env path: %s", self.env_path)
        self.installed = False
        self._is_attempted = False

        super(CondaPackage, self).__init__()

    @staticmethod
    def _string_from_start_of_json(string_with_json):
        # JSON can start with "{" or "["
        # via http://www.json.org/
        matches = re.compile("\{|\[").search(string_with_json)
        if matches:
            return string_with_json[matches.start():]
        else:
            _log.warn("Does not look like json: %s" % string_with_json)
            return None

    @property
    def _package_str(self):
        if self.version:
            ver_str = "{pkg}={ver}".format(
                pkg=self.package, ver=self.version.version_spec)
        else:
            ver_str = self.package
        return ver_str

    def is_attempted(self):
        return self._is_attempted

    def is_installed(self):
        return self.installed

    @property
    def bin_path(self):
        return os.path.join(self.env_path, "bin")

    def executable_path(self):
        return os.path.join(self.bin_path, self.executable)

    def apply_patches(self):
        for path, patch in self.patches:
            self._patch(path, patch)


    def _patch(self, path, patch):
        """Patch a path relative to conda root.

        The patch is relative to the tools/patches directory.
        """
        file_path = os.path.join(self.env_path, path)
        patch_path = os.path.join(
            util.file.get_project_path(), 'tools', 'patches', patch)
        subprocess.check_call(['patch', file_path, patch_path])

    @property
    def _package_installed(self):
        result = util.misc.run_and_print(["conda", "list", "-f", "-c", "-p", self.env_path, "--json", self.package], loglevel=logging.DEBUG, env=self.conda_env)
        if result.returncode == 0:
            command_output = result.stdout.decode("UTF-8")
            data = json.loads(self._string_from_start_of_json(command_output))
            if len(data) > 0:
                _log.debug('Conda package found: {}'.format(data))
                return True
        return False

    def verify_install(self):
        # if conda does not the package as installed, no need to check further
        installed_version = self.get_installed_version()
        if not installed_version:
            # report the package as not installed
            return False

        # if the package is installed, check the binary for executability
        if os.access(self.executable_path(), (os.X_OK | os.R_OK) if self.require_executability else os.R_OK):
            # optionally use the verify command, if specified
            if self.verifycmd:
                if os.system(self.verifycmd) == self.verifycode:
                    _log.debug("Validating with cmd: {}".format(self.verifycmd))
                    self.installed = installed_version
            else:
                self.installed = installed_version
        else:
            self.installed = False
        if self.installed:
            # call the post-verification command.
            # Useful for copying in license files, building databases, etc.
            # This is executed relative to the conda environment bin/
            # And has the active conda environment on the PATH
            if (not self.post_verify_cmd_executed) and self.post_verify_command:
                post_verify_command = shlex.split(self.post_verify_command)
                _log.debug("Running post-verification cmd: {}".format(self.post_verify_command))

                result = util.misc.run_and_print(post_verify_command, silent=False, check=False, env=self.conda_env, cwd=self.bin_path)
                post_verify_cmd_return_code = result.returncode
                if post_verify_cmd_return_code == self.post_verify_ret:
                    self.post_verify_cmd_executed = True
                else:
                    raise subprocess.CalledProcessError(post_verify_cmd_return_code, "Post verification command failed with exit %s: %s" % (post_verify_cmd_return_code, self.post_verify_command))

            return installed_version
        return False

    def _attempt_install(self):
        try:
            # check for presence of conda command
            util.misc.run_and_print(["conda", "-V"], buffered=True, check=True, env=self.conda_env, silent=True)
        except:
            _log.warning("conda NOT installed")
            self._is_attempted = True
            self.installed = False
            return
            #raise

        # conda-build is not needed for pre-built binaries from conda channels
        # though we may will need it in the future for custom local builds
        # try:
        #     util.misc.run_and_print(["conda", "build", "-V"], silent=True, check=True, env=self.conda_env)
        # except:
        #     _log.warning("conda-build must be installed; installing...")
        #     util.misc.run_and_print(["conda", "install", "-y", "conda-build"], check=True)

        # if the package is already installed, we need to check if the version is correct
        pkg_version = self.verify_install()
        if pkg_version:
            _log.debug("Currently installed version of {package}: {version}".format(
                package=self.package, version=pkg_version))
            # if the installed version is not the one specified
            if not pkg_version.satisfies(self.version):
                _log.debug("Expected version of {package}:            {version}".format(package=self.package, version=self.version))
                _log.debug("Incorrect version of {package} installed. Removing it...".format(package=self.package) )

                # uninstall the current (incorrect) version
                self.uninstall_package()
                # and continue to install...
            else:
                # if the package is installed and is the correct version
                # return so we don't bother installing
                return

        # install the package and verify
        _log.debug("Attempting install...")
        self.install_package()
        self.verify_install()
        self.post_install()

    def get_installed_version(self):
        # If we ever use conda to install pip packages as tools, "-c" needs to be removed
        run_cmd = ["conda", "list", "-c", "--json", "-f", "-p", self.env_path, self.package]

        result = util.misc.run_and_print(run_cmd, silent=True, check=False, env=self.conda_env)
        if result.returncode == 0:
            try:
                command_output = result.stdout.decode("UTF-8")
                data = json.loads(self._string_from_start_of_json(command_output))
            except:
                _log.warning("failed to decode JSON output from conda create: %s", result.stdout.decode("UTF-8"))
                return # return rather than raise so we can fall back to the next install method

            if data and len(data):
                if isinstance(data[0], dict):
                    installed_package_string = data[0]["dist_name"]
                else:
                    installed_package_string = data[0]
                package_info_re = re.compile(r"(?P<package_name>.*)-(?P<version>.*)-(?P<build_type>.*)")
                matches = package_info_re.match(installed_package_string)
                if matches:
                    installed_version = matches.group("version")
                    installed_package = matches.group("package_name")
                    installed_build_type = matches.group("build_type")
                    return CondaPackageVersion(installed_version, installed_build_type)
        return None

    def package_available(self):
        # If we ever use conda to install pip packages as tools, "-c" needs to be removed
        run_cmd = ["conda", "search", "--json", "-c", self.channel, self.package]

        result = util.misc.run_and_print(run_cmd, loglevel=logging.INFO, env=self.conda_env)
        if result.returncode == 0:
            try:
                command_output = result.stdout.decode("UTF-8")
                data = json.loads(self._string_from_start_of_json(command_output))
            except:
                _log.warning("failed to decode JSON output from conda search: %s", result.stdout.decode("UTF-8"))
                return # return rather than raise so we can fall back to the next install method

            if data and len(data):
                if self.package in data and "error" not in data:
                    for sub_pkg in data[self.package]:
                        if sub_pkg.get("version", "") == self.version.version_spec:
                            return True

        _log.info("Conda package for %s is not available on this platform.", self.package)
        return False

    def uninstall_package(self):
        run_cmd = ["conda", "remove", "-q", "-y", "--json", "-p", self.env_path, self.package]

        result = util.misc.run_and_print(
            run_cmd,
            loglevel=logging.INFO,
            env=self.conda_env)

        if result.returncode == 0:
            try:
                command_output = result.stdout.decode("UTF-8")
                data = json.loads(self._string_from_start_of_json(command_output))
            except:
                _log.warning("failed to decode JSON output from conda install: %s", result.stdout.decode("UTF-8"))
                return # return rather than raise so we can fall back to the next install method

            if data["success"] == True:
                _log.debug("Package removed.")
                #self.installed = False

        self.verify_install()

    def install_package(self):
        #if not self.package_available():
        #    _log.error("Conda package for %s cannot be installed; it is not available on this platform. Related functionality may not be available.", self.package)
        #    return

        # try to create the environment and install the package
        run_cmd = ["conda", "create", "-q", "-y", "--json", "-c", self.channel, "-p", self.env_path, self._package_str]

        _log.debug("Creating conda environment and installing package {package}".format(package=self._package_str))

        python_version = os.environ.get("TRAVIS_PYTHON_VERSION")
        if python_version:
            python_version = "python=" + python_version if python_version else ""
            run_cmd.extend([python_version])

        result = util.misc.run_and_print(run_cmd, loglevel=logging.INFO, env=self.conda_env)
        try:
            command_output = result.stdout.decode("UTF-8")
            data = json.loads(self._string_from_start_of_json(command_output))
        except:
            _log.warning("failed to decode JSON output from conda create: %s", result.stdout.decode("UTF-8"))
            return # return rather than raise so we can fall back to the next install method

        if "error" in data.keys() and "prefix already exists" in data["error"]:
            # the environment already exists
            # the package may not be installed...
            _log.debug("Conda environment already exists. Installing package...")

            result = util.misc.run_and_print(
                [
                    # --no-update-dependencies ensures subsequent installs do not bump versions of prior pinned installs
                    "conda", "install", "--json", "-c", self.channel, "-y", "-q", "--no-update-dependencies", "-p", self.env_path,
                    self._package_str
                ],
                loglevel=logging.DEBUG,
                env=self.conda_env,
            )

            if result.returncode == 0:
                try:
                    command_output = result.stdout.decode("UTF-8")
                    data = json.loads(self._string_from_start_of_json(command_output))
                except:
                    _log.warning("failed to decode JSON output from conda install: %s", result.stdout.decode("UTF-8"))
                    return # return rather than raise so we can fall back to the next install method

                if data["success"] == True:
                    _log.debug("Package installed.")
            self.apply_patches()

        else:
            if "success" in data.keys() and data["success"]:
                # we were able to create the environment and install the package
                _log.debug("Conda environment created and package installed.")

    def post_install(self):
        """
            Runs a shell command after package installation, 
            relative to the directory containing the executable
        """
        if self.post_install_command:
            destination_dir = os.path.dirname(os.path.realpath(self.executable_path()))
            return_code = os.system('cd "{}" && {}'.format(destination_dir, self.post_install_command))
            if self.post_install_ret is not None:
                assert return_code == self.post_install_ret


class DownloadPackage(InstallMethod):
    ''' This is an install method for downloading, unpacking, and post-
            processing straight from the source.
        target_rel_path is the executable's path relative to destination_dir
        destination_dir defaults to the project build directory
        post_download_command will be executed if it isn't None, in
            destination_dir.
        if post_download_ret != None, assert it is returned by
            post_download_command
    '''

    def __init__(
        self,
        url,
        target_rel_path,
        destination_dir=None,
        verifycmd=None,
        verifycode=0,
        require_executability=True,
        post_download_command=None,
        post_download_ret=0
    ):
        if destination_dir is None:
            destination_dir = util.file.get_build_path()
        self.url = url
        self.targetpath = os.path.join(destination_dir, target_rel_path)
        self.destination_dir = destination_dir
        self.verifycmd = verifycmd
        self.verifycode = verifycode
        self.installed = False
        self.require_executability = require_executability
        self.post_download_command = post_download_command
        self.post_download_ret = post_download_ret
        self.download_file = None
        InstallMethod.__init__(self)

    def is_installed(self):
        return self.installed

    def executable_path(self):
        return self.installed and self.targetpath or None

    def verify_install(self):
        if os.access(self.targetpath, (os.X_OK | os.R_OK) if self.require_executability else os.R_OK):
            if self.verifycmd:
                _log.debug("validating")
                self.installed = (os.system(self.verifycmd) == self.verifycode)
            else:
                self.installed = True
        else:
            self.installed = False
        return self.installed

    def _attempt_install(self):
        if not self.verify_install():
            self.pre_download()
            self.download()
            self.post_download()
            self.verify_install()

    def pre_download(self):
        pass

    def download(self):
        download_dir = tempfile.gettempdir()
        util.file.mkdir_p(download_dir)
        filepath = urlparse(self.url).path
        file_basename = filepath.split('/')[-1]
        _log.info("Downloading from %s to %s/%s ...", self.url, download_dir, file_basename)
        urlretrieve(self.url, os.path.join(download_dir, file_basename))
        self.download_file = file_basename
        self.unpack(download_dir)

    def post_download(self):
        if self.post_download_command:
            return_code = os.system('cd "{}" && {}'.format(self.destination_dir, self.post_download_command))
            if self.post_download_ret is not None:
                assert return_code == self.post_download_ret

    def unpack(self, download_dir):
        _log.debug("unpacking")
        util.file.mkdir_p(self.destination_dir)
        if self.download_file.endswith('.zip'):
            if os.system("unzip -o %s/%s -d %s > /dev/null" % (download_dir, self.download_file, self.destination_dir)):

                return
            else:
                os.unlink(os.path.join(download_dir, self.download_file))
        elif (
            self.download_file.endswith('.tar.gz') or self.download_file.endswith('.tgz') or
            self.download_file.endswith('.tar.bz2') or self.download_file.endswith('.tar')
        ):
            if self.download_file.endswith('.tar'):
                compression_option = ''
            elif self.download_file.endswith('.tar.bz2'):
                compression_option = 'j'
            else:
                compression_option = 'z'
            untar_cmd = "tar -C {} -x{}pf {}/{}".format(
                self.destination_dir, compression_option, download_dir, self.download_file
            )
            _log.debug("Untaring with command: %s", untar_cmd)
            exitCode = os.system(untar_cmd)
            if exitCode:
                _log.info("tar returned non-zero exitcode %s", exitCode)
                return
            else:
                _log.debug("tar returned with exit code 0")
                os.unlink(os.path.join(download_dir, self.download_file))
        else:
            shutil.move(
                os.path.join(download_dir, self.download_file), os.path.join(self.destination_dir, self.download_file)
            )