Spintrum/setup.py

from distutils.command.build import build as _build
import os
import sys
import shutil
import errno
import tempfile
import subprocess
import glob
import re
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 get_version():
    VERSIONFILE = os.path.join(os.path.dirname(os.path.abspath(__file__)),"spintrum/meta.py")
    verstrline = open(VERSIONFILE, "rt").read()
    VSRE = r"^__version__ = ['\"]([^'\"]*)['\"]"
    mo = re.search(VSRE, verstrline, re.M)
    if mo:
        return str(mo.group(1))
    else:
        raise RuntimeError("Unable to find version string in %s." % (VERSIONFILE,))

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()
    sys.stderr.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()
    sys.stderr.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 list(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()
    openblas_target_optimization = "-Ofast"
    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.")

        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()
        spintrum_optimization_target = "-Ofast"
        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=get_version(),
      description='Software for spin systems simulation',
      url='http://www.afach.de/',
      author='Samer Afach',
      author_email='samer@afach.de',
      license='Dual',
      packages=['spintrum'],
      package_data={'spintrum': [lib_file] + recursive_glob("examples")},
      install_requires=['numpy', 'mpmath'],
      python_requires = '>=3.3',
      include_package_data=True, # enable including files with MANIFEST.in
      zip_safe=False,
      cmdclass=dict(build=DependenciesBuilder)
      )