setup.py

#!/usr/bin/env python

import os
import shutil
import glob
import sys
import shlex
import sysconfig
import subprocess

try:
    import multiprocessing
    _cpu_count = multiprocessing.cpu_count()
except Exception:
    _cpu_count = 1

import setuptools
from setuptools import distutils
from setuptools import Command
from setuptools import setup
from setuptools import find_packages
from setuptools.dist import Distribution as _Distribution
from setuptools.command.install import install as _install

from distutils.command.build import build as _build
from distutils import log

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


class cmake(Command):

    description = "Runs cmake configuration"

    user_options = [
        ('cmake-command', None, 'CMake program path'),
        ('make-command', None, 'GNU Make program path'),
        ('source-path', None, 'Path to the source tree containing CMakeLists.txt'),
        ('staging-path', None, 'Path to the install staging directory.'),
        ('extra-args', None, 'Extra arguments that will be passed to cmake'),
        ('extra-envs', None, 'Extra environment variables that will be used to build'),
        ('force', 'f', 'forcibly build everything (ignore file timestamps)'),
        ('debug', 'g', 'compile with debugging'),
        ('serial', 's' 'force serial build')
    ]

    _cmake_config = {
        'PYTHON_EXECUTABLE': sys.executable,
        'BUILD_TESTING': False,
        'BUILD_SHARED_LIBS': True,
        'BUILD_DOCUMENTATION': False,
        'VTK_WRAP_PYTHON': True
    }

    def initialize_options(self):

        self.cmake_command = None
        self.make_command = None
        self.source_path = None
        self.build_path = None
        self.staging_path = None
        self.cmake_config = None
        self.extra_args = None
        self.extra_envs = None
        self.force = None
        self.debug = None
        self.serial = None
        self.deploy_path = None

    def finalize_options(self):

        self.set_undefined_options('clone', ('clone_path', 'source_path'))

        build = self.distribution.get_command_obj('build')
        build.ensure_finalized()

        install = self.distribution.get_command_obj('install')
        install.ensure_finalized()

        if self.deploy_path is None:
            self.deploy_path = build.build_base

        if self.debug:
            self._build_type = 'Debug'
        else:
            self._build_type = 'Release'

        if self.cmake_command is None:
            self.cmake_command = 'cmake'

        if self.make_command is None:
            self.make_command = 'make'

        if self.build_path is None:
            self.build_path = 'cmake_build'

        if self.staging_path is None:
            self.staging_path = 'staging'

        if self.cmake_config is None:
            self.cmake_config = {}
        self.cmake_config.update(cmake._cmake_config)
        if 'CMAKE_BUILD_TYPE' not in self.cmake_config:
            self.cmake_config['CMAKE_BUILD_TYPE'] = self._build_type
        self.cmake_config['CMAKE_INSTALL_PREFIX'] = os.path.abspath(self.staging_path)

        # need to get the actual install path
        self.cmake_config['CMAKE_INSTALL_RPATH'] = os.path.abspath(
            os.path.join(install.install_lib, 'vtk')
        )
        self.cmake_config['PYTHON_INCLUDE_DIR'] = sysconfig.get_path('include')
        pylib = self.get_python_library()
        if pylib:
            self.cmake_config['PYTHON_LIBRARY'] = pylib

        if self.extra_args is None:
            self.extra_args = ''

        if self.extra_envs is None:
            self.extra_envs = ''

    @classmethod
    def get_python_library(cls):
        '''
        Return the /path/to/libpython*.so or similar python library of the current
        python.  There doesn't seem to be a general place to find this, so we have
        to hack it out.
        '''

        # this seems to work with mac homebrew and ubuntu, more customizations will
        # probably be necessary
        libdir = sysconfig.get_config_var('LIBPL')

        libglob = glob.glob(os.path.join(libdir, 'libpython*'))
        for lib in libglob:
            if os.path.splitext(lib)[1] in ('.so', '.dll', '.dylib'):
                return lib
        return None

    @classmethod
    def config_to_cmake_args(cls, config):
        '''
        Convert a configuration dictionary into a cmake argument list.
        '''

        args = []
        for key, value in config.iteritems():
            if not isinstance(key, basestring) or \
               not isinstance(value, (basestring, bool)):
                raise Exception("Invalid configuration object")
            if isinstance(value, basestring):
                spec = '%s:STRING=%s' % (key, value)
            elif isinstance(value, bool):
                if value:
                    value = 'ON'
                else:
                    value = 'OFF'
                spec = '%s:BOOL=%s' % (key, value)

            args.extend(['-D', spec])
        return args

    def run(self):
        '''
        Run cmake, make, and make install.  Could be split into subcommands in the future.
        '''
        cmake_args = [self.cmake_command]
        cmake_args += self.config_to_cmake_args(self.cmake_config)
        cmake_args += shlex.split(self.extra_args)
        cmake_args += [self.source_path]

        make_command = self.make_command
        if not self.serial:
            make_command += ' -j%i' % _cpu_count

        env = [value.split('=', 1) for value in shlex.split(self.extra_envs)]
        env = {value[0].strip(): value[1] for value in env if len(value) == 2}

        build_stamp = os.path.join(self.build_path, '.built')
        install_stamp = os.path.join(self.staging_path, '.installed')

        def conf():
            self.ensure_dirname('source_path')
            self.ensure_dirname('build_path')
            log.info(' '.join(cmake_args))
            subprocess.check_call(
                ' '.join(cmake_args),
                cwd=self.build_path,
                shell=True,
                env=env
            )

        def build():
            self.ensure_dirname('source_path')
            self.ensure_dirname('build_path')
            log.info(make_command)
            subprocess.check_call(
                make_command,
                cwd=self.build_path,
                shell=True,
                env=env
            )
            subprocess.call(
                ' '.join([self.cmake_command, '-E', 'touch', build_stamp]),
                shell=True
            )

        def install():
            self.ensure_dirname('build_path')
            log.info('make install')
            subprocess.check_call(
                ' '.join([self.make_command, 'install']),
                cwd=self.build_path,
                shell=True,
                env=env
            )
            subprocess.call(
                ' '.join([self.cmake_command, '-E', 'touch', install_stamp]),
                shell=True
            )

        # make the build directory if necessary
        try:
            self.mkpath(self.build_path)
        except Exception:
            pass

        self.make_file(
            os.path.join(self.source_path, 'CMakeLists.txt'),
            os.path.join(self.build_path, 'CMakeCache.txt'),
            conf, [],
            exec_msg='Running cmake configure...',
            skip_msg='Skipping configure step, use -f to force.'
        )

        self.make_file(
            os.path.join(self.build_path, 'CMakeCache.txt'),
            build_stamp,
            build, [],
            exec_msg='Running make...',
            skip_msg='Skipping build step, use -f to force.'
        )

        self.make_file(
            build_stamp,
            install_stamp,
            install, [],
            exec_msg='Running make install to staging path...',
            skip_msg='Skipping installation to staging path, use -f to force.'
        )

        # get a list of scripts/binaries to install
        # (maybe later add to install_scripts)
        binaries = glob.glob(os.path.join(self.staging_path, 'bin', '*'))

        # get the path to the main vtk package
        vtkpath = glob.glob(
            os.path.join(
                self.staging_path,
                'lib*',
                'python*',
                'site-packages'
            )
        )
        if not vtkpath:
            raise Exception(
                "Could not find main vtk package directory in %s" % self.staging_path
            )
        vtkpath = vtkpath[0]

        # copy files to the build tree
        try:
            shutil.rmtree(os.path.join(self.deploy_path, 'lib', 'vtk'))
        except Exception:
            pass
        shutil.copytree(os.path.join(vtkpath, 'vtk'), os.path.join(self.deploy_path, 'lib', 'vtk'))
        for f in glob.glob(os.path.join(self.staging_path, 'lib*', '*')):
            if os.path.islink(f):
                os.symlink(
                    os.readlink(f),
                    os.path.join(self.deploy_path, 'lib', 'vtk', os.path.basename(f))
                )
            elif os.path.isfile(f):
                shutil.copy2(f, os.path.join(self.deploy_path, 'lib', 'vtk', os.path.basename(f)))


class clone(Command):

    description = "Downloads a cmake project source from git."

    user_options = [
        ('git-command', None, 'Git program path'),
        ('repository', 'r', 'Source repository url'),
        ('head', None, 'Tag or branch to clone'),
        ('clone-path', None, 'The path to clone into')
    ]

    def initialize_options(self):

        self.git_command = None
        self.repository = None
        self.head = None
        self.clone_path = None

    def finalize_options(self):

        build = self.distribution.get_command_obj('build')
        build.ensure_finalized()

        if self.git_command is None:
            self.git_command = 'git'

        if self.repository is None:
            self.repository = 'https://github.com/Kitware/VTK.git'

        if self.head is None:
            self.head = 'master'

        if self.clone_path is None:
            self.clone_path = 'vtk-source'
        self.clone_path = os.path.abspath(self.clone_path)

    def run(self):

        def exe():
            self.spawn(
                [
                    self.git_command,
                    'clone',
                    '--depth', '1',
                    '--branch', self.head,
                    '--single-branch',
                    self.repository,
                    self.clone_path
                ]
            )
        self.make_file(
            [],
            os.path.join(self.clone_path, '.git', 'config'),
            exe, [],
            exec_msg='Cloning VTK repository',
            skip_msg='Repository already cloned, skipping'
        )


class build(_build):
    sub_commands = _build.sub_commands + [('clone', None), ('cmake', None)]


class Distribution(_Distribution):

    def _include_package_dir(self, obj):
        if self.package_dir is None:
            self.package_dir = {}
        self.package_dir.update(obj)
        self.reinitialize_command('install', reinit_subcommands=1)
        install = self.get_command_obj('install')
        install.skip_build = True
        install.ensure_finalized()

    def has_ext_modules(self):
        return True


setup(
    name='VTK python bindings installer',
    version='6.1',
    author='Kitware, Inc.',
    zip_safe=False,
    cmdclass={
        'cmake': cmake,
        'build': build,
        'clone': clone
    },
    platforms=[distutils.util.get_platform()],
    distclass=Distribution
)