Spintrum/setup.py

from setuptools import setup, find_packages
from distutils.command.build import build as _build
from spintrum import meta
import os
import sys
import shutil
import errno
import tempfile
import subprocess
import glob
import fnmatch

try:
    from setuptools import setup, Extension
except ImportError:
    from distutils.core import setup
    from distutils.extension import Extension

del os.link #solve hardlinking problem when using NTFS drive

script_dir = os.path.dirname(os.path.realpath(__file__))
start_working_dir = os.getcwd()
temp_dir = tempfile.gettempdir()


def _print_error(msg):
    sys.stderr.write("\n\n" + error_color_start + "ERROR: " + msg + color_end + " \n\n\n ")
    sys.stderr.flush()


def which(program):
    """
    Checks if a program exists, and returns its path
    :param program: global program name
    :return: program path if it exists, otherwise None
    """
    import os
    def is_exe(fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    fpath, fname = os.path.split(program)
    if fpath:
        if is_exe(program):
            return program
    else:
        for path in os.environ["PATH"].split(os.pathsep):
            path = path.strip('"')
            exe_file = os.path.join(path, program)
            if is_exe(exe_file):
                return exe_file

    return None


def check_programs_availability(list_of_programs):
    for prog in list_of_programs:
        if which(prog) is None:
            _print_error("The required program " + prog + " was not found in the global scope (in PATH). Please install it. Exiting...")
            sys.exit(1)
    print("All required programs were found to be installed. Proceeding with building and installation.")


def mkdir_p(path):
    try:
        os.makedirs(path)
    except OSError as exc:
        if exc.errno == errno.EEXIST and os.path.isdir(path):
            pass
        else:
            raise


def deal_with_error_code(err_code, operation_name):
    if err_code != 0:
        _print_error("A non-zero error code was returned at operation: " +
                         str(operation_name) + ". Code returned is: " +
                         str(err_code) + ". Exiting!")
        sys.exit(1)


def inplace_change(filename, old_string, new_string):
    # Safely read the input filename using 'with'
    with open(filename) as f:
        s = f.read()
        if old_string not in s:
            print('"{old_string}" not found in {filename}.'.format(**locals()))
            return

    # Safely write the changed content, if found in the file
    with open(filename, 'w') as f:
        print('Changing "{old_string}" to "{new_string}" in {filename}'.format(**locals()))
        s = s.replace(old_string, new_string)
        f.write(s)


def clone_git_repository(dir, repos_link):
    try:
        shutil.rmtree(dir)
    except FileNotFoundError:
        pass
    mkdir_p(dir)
    print("Cloning repository: " + repos_link + "...")
    err_code = subprocess.call(["git clone " + repos_link + " " + dir], shell=True)
    deal_with_error_code(err_code, "Cloning: " + repos_link)
    print("Done cloning repository: " + repos_link)

question_color_start = "\x1b[0;37;44m"
error_color_start = "\x1b[0;37;41m"
color_end = "\x1b[0m"

def ask_openblas_optimization_level():
    print("\n")
    print(question_color_start +
          "Which optimization option would you like to use to compile OpenBLAS?" +
          color_end)
    print("For more information on what the options mean, please visit:")
    print("https://gcc.gnu.org/onlinedocs/gcc/Optimize-Options.html")
    print("1) -O2")
    print("2) -O3")
    print("3) -Ofast")
    print("4) Keep the default")
    print("5) Skip compiling OpenBLAS and assume it exists in the system")
    print(question_color_start + "Please enter the option number: " + color_end)
    sys.stdout.flush()
    option_chosen = input()
    if option_chosen.replace(" ", "") == '1':
        return "-O2"
    elif option_chosen.replace(" ", "") == '2':
        return "-O3"
    elif option_chosen.replace(" ", "") == '3':
        return "-Ofast"
    elif option_chosen.replace(" ", "") == '4':
        return ""
    elif option_chosen.replace(" ", "") == '5':
        return "SKIPOPENBLAS"
    else:
        raise IndexError(error_color_start + "Error: An unknown option was entered" + color_end)


def ask_spintrum_optimization_level():
    print("\n")
    print(question_color_start +
          "Which optimization option would you like to use to compile Spintrum?" +
          color_end)
    print("For more information on what the options mean, please visit:")
    print("https://gcc.gnu.org/onlinedocs/gcc/Optimize-Options.html")
    print("1) -g")
    print("2) -O1")
    print("3) -O2")
    print("4) -O3")
    print("5) -Ofast")
    print(question_color_start + "Please enter the option number: " + color_end)
    sys.stdout.flush()
    option_chosen = input("")
    if option_chosen.replace(" ", "") == '1':
        return "-g"
    elif option_chosen.replace(" ", "") == '2':
        return "-O1"
    elif option_chosen.replace(" ", "") == '3':
        return "-O2"
    elif option_chosen.replace(" ", "") == '4':
        return "-O3"
    elif option_chosen.replace(" ", "") == '5':
        return "-Ofast"
    else:
        raise IndexError(error_color_start + "Error: An unknown option was entered" + color_end)


def recursive_glob(rootdir, pattern='*'):
    lst = [os.path.join(looproot, filename)
            for looproot, _, filenames in os.walk(rootdir)
            for filename in filenames
            if fnmatch.fnmatch(filename, pattern)]
    return lst


def get_openblas():
    openblas_dir = "3rdparty/OpenBLAS"
    openblas_dir_install = "spintrum/OpenBLAS_install"
    openblas_full_path = os.path.join(start_working_dir, openblas_dir)
    openblas_full_path_install = os.path.join(start_working_dir, openblas_dir_install)
    openblas_git = "https://github.com/xianyi/OpenBLAS.git"

    openblas_target_optimization = ask_openblas_optimization_level()
    if openblas_target_optimization != "SKIPOPENBLAS":
        clone_git_repository(openblas_full_path, openblas_git)
        if openblas_target_optimization != "":
            if not os.path.isdir(openblas_dir):
                _print_error("Cannot open expected OpenBLAS directory " + openblas_dir)
                sys.exit(1)

            makefiles_to_glob = os.path.join(openblas_dir, "Makefile*")
            makefiles = glob.glob(makefiles_to_glob)

            print("Number of files found: " + str(len(makefiles)))

            for f in makefiles:
                print(f)
                inplace_change(f, "-O2", openblas_target_optimization)

        # make openblas
        os.chdir(openblas_dir)
        make_openblas_err_code = subprocess.call(["make"], shell=True)
        deal_with_error_code(make_openblas_err_code, "Making OpenBLAS")

        # install openblas
        install_openblas_err_code = subprocess.call(
            ["make PREFIX=" + openblas_full_path_install +
             " install"], shell=True)
        deal_with_error_code(install_openblas_err_code, "Installing OpenBLAS")
        os.chdir(start_working_dir)  # restore working dir


dir_name = "spintrum"
lib_file = "lib/libspintrum.so"


class DependenciesBuilder(_build):
    def run(self):
        sys.stdout.flush()
        sys.stderr.flush()
        if sys.platform.startswith('win'):
            _print_error("You cannot build Spintrum on Windows. It is not supported.")

        python_version = float(str(sys.version_info[0]) + "." + str(sys.version_info[1]))
        if python_version < 3.3:
            _print_error("You must have python version >= 3.3 to use Spintrum")

        check_programs_availability(["cmake","make","g++","gcc","gfortran","git","python3"])

        print("Building Spintrum C/C++ extension prerequisites")

        get_openblas()

        ############## get Polymath
        polymath_git = "https://git.afach.de/samerafach/Polymath"
        polymath_dir = "3rdparty/Polymath"
        polymath_full_path = os.path.join(start_working_dir, polymath_dir)
        clone_git_repository(polymath_full_path, polymath_git)

        ############## build spintrum
        spintrum_optimization_target = ask_spintrum_optimization_level()
        print("Building Spintrum extension...")
        err_code = subprocess.call(["cmake . -DPythonInstallation=1 -DOptimizationLevel=" + spintrum_optimization_target], shell=True)
        deal_with_error_code(err_code, "CMaking spintrum")
        err_code = subprocess.call(["make"], shell=True)
        deal_with_error_code(err_code, "Making spintrum")
        print("Done building spintrum extension.")


setup(name='spintrum',
      version=meta.__version__,
      description='Software for spin systems simulation',
      url='http://www.afach.de/',
      author='Samer Afach',
      author_email='samer@afach.de',
      license='MPL',
      packages=['spintrum'],
      include_package_data=True, # enable including files with MANIFEST.in
      package_data={'': [lib_file]},
      zip_safe=False,
      cmdclass=dict(build=DependenciesBuilder)
      )