bundler.py

#!/usr/bin/env python
"""
Usage:

add this to your setup.py:
...
      cmdclass={'bundle': BundleInstall},
...

Will add a bundle setup.py command that create a bundle archive in dist directory by running
$ python setup.py bundle

"""
import os
import re
import subprocess
import sys
import tempfile
import shutil
from distutils.core import Command


def fix_file(path, sdesc):
    tmpfile = path + '.new'
    os.rename(path, tmpfile)
    with open(tmpfile, 'w+') as tdesc:
        tdesc.write('# !/bin/sh\n')
        tdesc.write('"exec" "`dirname $0`/python" "$0" "$@"\n')
        tdesc.write(sdesc.read())
    sdesc.close()
    os.rename(tmpfile, path)


def fix_env_shebang(env):
    bin_path = os.path.join(env, 'bin')
    for binary in os.listdir(bin_path):
        binary_path = os.path.join(bin_path, binary)
        if not os.path.isfile(binary_path):
            continue

        with open(binary_path, 'r') as fdesc:
            try:
                firstline = fdesc.readline()
                if re.match('#!.+bin/python.*', firstline):
                    fix_file(binary_path, fdesc)
                    print("Fixing %s" % binary_path)
            except UnicodeDecodeError:
                continue


def create_conda_env(path):
    command = ' '.join(['conda', 'create', '--yes', '--copy', '--prefix', path, 'python'])
    print("Running: %s" % command)
    p = subprocess.Popen(command, stdout=sys.stdout, stderr=sys.stderr, shell=True)
    p.communicate()
    p.wait()
    if (p.returncode != 0):
        raise RuntimeError("Failed to create conda env in %s" % path)


def compress_bundle(path, archive_name):
    compress_command = "tar -C {parent_dir} -zcf {archive_name} {path}".format(
        parent_dir=os.path.dirname(path),
        path=os.path.basename(path),
        archive_name=archive_name)
    print("Running: %s" % compress_command)
    subprocess.check_call(compress_command, shell=True)


def install_package(path, setup_file, *args):
    python_path = os.path.join(path, "bin", "python")
    command = [python_path, setup_file, "install"] + list(args)
    print("Running: %s" % command)
    subprocess.check_call(command)


class BundleInstall(Command):
    user_options = []

    def __init__(self, dist):
        super().__init__(dist)

    def _get_output_build_name(self):
        return "{name}-{version}".format(
            name=self.distribution.metadata.get_name(),
            version=self.distribution.metadata.get_version()
        )

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def run(self):
        build_dir = tempfile.mkdtemp()
        env_path = os.path.join(build_dir, self._get_output_build_name())
        tmp_archive_name = "{env_path}-bundle.tar.gz".format(env_path=env_path)
        dist_directory = os.path.join(os.path.dirname(self.distribution.script_name),
                                      'dist')
        dist_file = os.path.join(dist_directory, os.path.basename(tmp_archive_name))

        create_conda_env(env_path)
        fix_env_shebang(env_path)
        install_package(env_path, self.distribution.script_name)
        fix_env_shebang(env_path)
        compress_bundle(env_path, tmp_archive_name)

        if not os.path.isdir(dist_directory):
            os.makedirs(dist_directory)
        print("Renaming %s -> %s" % (tmp_archive_name, dist_file))
        shutil.move(tmp_archive_name, dist_file)