snap.py 10.4 KB
#!/usr/bin/python3

"""Create snaps from Python programs."""

import os
import sys
import glob
import json
import yaml
import shutil
import tempfile
import subprocess

from argparse import ArgumentParser
from configparser import ConfigParser
from contextlib import ExitStack, contextmanager
from email import message_from_bytes
from email.utils import formataddr
from urllib.request import urlopen


@contextmanager
def chdir(new_dir):
    cwd = os.getcwd()
    try:
        os.chdir(new_dir)
        yield
    finally:
        os.chdir(cwd)


class Directories:
    def __init__(self, base=None):
        self.base = base
        self.bin = None
        self.meta = None
        self.tmp = None
        self._resources = ExitStack()

    def __enter__(self):
        self.tmp = tempfile.mkdtemp()
        self._resources.callback(shutil.rmtree, self.tmp)
        if self.base is None:
            self.base = tempfile.mkdtemp()
            self._resources.callback(shutil.rmtree, self.base)
        if not os.path.isdir(self.base):
            print('Missing base directory:', self.base, file=sys.stderr)
            sys.exit(1)
        self.meta = os.path.join(self.base, 'meta')
        self.bin = os.path.join(self.base, 'bin')
        os.makedirs(self.meta, exist_ok=True)
        os.makedirs(self.bin, exist_ok=True)
        return self

    def __exit__(self, *exc):
        self._resources.close()
        self.base = None
        self.bin = None
        self.tmp = None
        return False


class Package:
    def __init__(self, name, dirs):
        self.name = name
        self.vendor = None
        self.version = None
        self.source = None
        self.description = None
        self.dirs = dirs

    def _pexcmd(self, config, binary):
        raise NotImplementedError

    def _entry_point(self, config):
        return config['project']['entry_point']

    def write_yaml(self, meta_dir):
        yaml_path = os.path.join(meta_dir, 'package.yaml')
        yaml_dict = dict(
            name=self.name,
            vendor=self.vendor,
            version=self.version,
            source=self.source,
            )
        binaries = [dict(
            name='bin/{}'.format(self.name),
            description=self.description.splitlines()[0],
            )]
        yaml_dict['binaries'] = binaries
        with open(yaml_path, 'w', encoding='utf-8') as fp:
            yaml.dump(yaml_dict, fp, default_flow_style=False)

    def write_readme(self, meta_dir):
        readme_path = os.path.join(meta_dir, 'readme.md')
        with open(readme_path, 'w', encoding='utf-8') as fp:
            print(self.description, file=fp)

    def make_pex(self, config, bin_dir):
        # Now, run pex to create the zip file executable.  XXX Currently pex
        # requires an explicit entry point be specified.  This is usually
        # defined in the package's metadata, but that's not available until we
        # download the package and extract its egg-info.  For now, don't do the
        # double download (one to get the entry point and one by pex).  Also
        # see https://github.com/pantsbuild/pex/issues/23
        #
        # For now, require the entry point to be specified in the .ini file.
        binary = '{bin_dir}/{name}.pex'.format(
            bin_dir=bin_dir,
            name=self.name,
            )
        subprocess.check_call(self._pexcmd(config, binary), shell=True)
        unpexed = os.path.splitext(binary)[0]
        self.tweak_pex_file(binary, unpexed)
        os.chmod(unpexed, 0o775)
        os.remove(binary)

    def tweak_pex_file(self, binary, unpexed):
        # https://github.com/pantsbuild/pex/issues/53
        # Fix the shebang line.
        with open(binary, 'rb') as src:
            with open(unpexed, 'wb') as dst:
                skipped_shebang = False
                while True:
                    chunk = src.read(4096)
                    if len(chunk) == 0:
                        break
                    if skipped_shebang:
                        dst.write(chunk)
                        continue
                    if chunk[0:2] != b'#!':
                        dst.write(chunk)
                        skipped_shebang = True
                        continue
                    # Search for the \n and discard everything before it.
                    try:
                        i = chunk.index(b'\n')
                    except ValueError:
                        # No newline in text.
                        dst.write(chunk)
                        skipped_shebang = True
                        continue
                    dst.write(b'#!/usr/bin/python3\n')
                    dst.write(chunk[i+1:])
                    skipped_shebang = True


class PyPIPackage(Package):
    def __init__(self, name, dirs):
        super().__init__(name, dirs)
        # Get the PyPI metadata.
        with urlopen('http://pypi.python.org/pypi/{}/json'.format(name)) as f:
            metadata = json.loads(f.read().decode('utf-8'))
        info = metadata['info']
        self.vendor = formataddr((info['author'], info['author_email']))
        self.version = info['version']
        self.source = info['home_page']
        self.description = info['description']

    def _entry_point(self, config):
        try:
            return super()._entry_point(config)
        except KeyError:
            print('No ini file entry point; searching red-dove.com')
            # Try to use an external service to find the entry point.  The
            # other possibility is to pre-download the package adn get the
            # egg-info out of that, but for now let's not do double
            # downloads.
            url = ('http://www.red-dove.com/pypi/projects/{}/{}/'
                   'package-{}.json'.format(
                       self.name[0].upper(),
                       self.name,
                       self.version,
                       ))
            with urlopen(url) as f:
                metadata = json.loads(f.read().decode('utf-8'))
            eps = metadata['exports']['scripts']['console']
            assert len(eps) == 1, eps
            return eps[0].rsplit('=')[-1].strip()

    def _pexcmd(self, config, binary):
        ep = self._entry_point(config)
        print('Using entry point:', ep)
        return 'pex {verbose} -o {binary} -r {name} -e {entry_point}'.format(
            verbose=('-v' if config['pex']['verbose'] else ''),
            binary=binary,
            name=self.name,
            entry_point=ep,
            )


class GitPackage(Package):
    def __init__(self, name, url, dirs):
        super().__init__(name, dirs)
        self.egg_info = None
        self.dirs.wheels = os.path.join(self.dirs.tmp, 'wheels')
        os.makedirs(self.dirs.wheels, exist_ok=True)
        # The package isn't on PyPI, so let's clone the repo and try to
        # extract the metadata from there.
        with chdir(self.dirs.tmp):
            subprocess.check_call(
                'git clone {} upstream'.format(url), shell=True)
            os.chdir('upstream')
            subprocess.check_call('python3 setup.py bdist_wheel --universal',
                                  shell=True)
            # Try to extract various bits of information from the egg-info.
            egg_infos = glob.glob('*.egg-info')
            assert len(egg_infos) == 1, egg_infos
            self.egg_info = os.path.abspath(egg_infos[0])
            with open(os.path.join(self.egg_info, 'PKG-INFO'), 'rb') as fp:
                info = message_from_bytes(fp.read())
            # Save the wheel file before the tempdir gets deleted.  There
            # should probably be a better place for this.
            wheels = glob.glob('dist/*.whl')
            assert len(wheels) == 1, wheels
            shutil.copy(wheels[0], self.dirs.wheels)
        self.vendor = formataddr((info['author'], info['author-email']))
        self.version = info['version']
        self.source = info['home-page']
        self.description = (info['summary']
                            if info['description'] == 'UNKNOWN'
                            else info['description'])

    def _entry_point(self, config):
        # Allow any entry point defined in the .ini file to override the
        # package's entry points.  If there is none defined in the ini file,
        # grab it out of the egg info.
        try:
            return super()._entry_point(config)
        except KeyError:
            ep_txt = os.path.join(self.egg_info, 'entry_points.txt')
            eps = ConfigParser()
            eps.read(ep_txt)
            assert len(eps['console_scripts']) == 1, eps['console_scripts']
            return list(eps['console_scripts'].values())[0]

    def _pexcmd(self, config, binary):
        ep = self._entry_point(config)
        print('Using entry point:', ep)
        return ('pex {verbose} -o {binary} -r {name} --no-pypi '
                '--repo={wheel_dir} -e {entry_point}'.format(
                    verbose=('-v' if config['pex']['verbose'] else ''),
                    binary=binary,
                    name=self.name,
                    entry_point=ep,
                    wheel_dir=self.dirs.wheels,
                    ))


def main(argv=None):
    if argv is None:
        argv = sys.argv[1:]
    parser = ArgumentParser(
        prog='snap',
        description='Create a snap from a Python program')
    parser.add_argument('config', metavar='config file', nargs=1,
                        help='ini file describing snap')
    parser.add_argument('-b', '--base',
                        help='Base directory for snappy layout')
    args = parser.parse_args(argv)
    config = ConfigParser()
    config.read(args.config)
    # Start by getting the name of the PyPI package.
    project_info = config['project']
    name = project_info['name']
    try:
        origin = project_info['origin']
    except KeyError:
        origin = None
    with Directories(args.base) as dirs:
        if origin is None:
            package = PyPIPackage(name, dirs)
        else:
            protocol, url = origin.split()
            if protocol != 'git':
                raise NotImplementedError('Protocol not supported: {}'.format(
                    protocol))
            package = GitPackage(name, url, dirs)
        package.write_yaml(dirs.meta)
        package.write_readme(dirs.meta)
        package.make_pex(config, dirs.bin)
        # This leaves the snap in the current working directory
        subprocess.check_call('snappy build {}'.format(dirs.base), shell=True)
    return 0


if __name__ == '__main__':
    sys.exit(main())