Create a self-compressed Python script

containment.py

'''Python self-contained module maker

Creates a small Python script containing the neccessary complete
modules in a tar.bz2 archive. The generated script unpacks this
and it can be used as a stand-alone script without explicitly
installing the whole module.'''

import os
import sys
import os.path
import stat
import time

from hashlib import sha256
from tarfile import *
from base64 import b64encode

try:
    from cStringIO import StringIO

except ImportError:
    from StringIO import StringIO

class Compressed(object):
    '''Compressor object

    Members:
    _paths -- List of paths to recursively scan for files
    _dependencies -- List of dependencies to check after extract.
    _data -- Raw data
    _checksum -- sha256 checksum of _data
    '''

    def __init__(self):
        self._paths = []
        self._dependencies = []
        self._data = None
        self._checksum = None

    def add_path(self, path, modname):
        '''Add a path to the directory scanlist

        Arguments:
        path -- Path to scan
        modname -- Add path as module modname
        '''

        self._paths.append((path, modname))
        self._data = None
        self._checksum = None

    def add_module(self, modname, path = None):
        '''Add a python module to the directory scanlist

        Arguments:
        modname -- Python module name
        path -- Give the path of the module explicitly
        '''

        oldpaths = sys.path
        try:
            if path is not None:
                sys.path = [ path ]

            mod = __import__(modname)

        finally:
            sys.path = oldpaths

        (dirname, filename) = os.path.split(mod.__file__)
        (root, ext) = os.path.splitext(filename)

        pyfile = os.path.join(dirname, '%s.py' % os.path.join(dirname, root))

        if not os.path.isfile(pyfile):
            raise Exception("Missing module .py file %s" % pyfile)

        if root == '__init__':
            path = dirname

        else:
            path = pyfile

        self.add_path(path, modname)

    def add_dependency(self, modname):
        '''Add a module dependency'''

        self._dependencies.append(modname)

    @property
    def data(self):
        '''Get the raw compressed data'''

        if self._data is None:
            self.__compress()

        return self._data

    @property
    def checksum(self):
        '''Get the sha256 checksum'''

        if self._checksum is None:
            self._checksum = sha256(self.data).hexdigest()

        return self._checksum

    def exclude_file(self, filename):
        '''Exclude pyc files. Callback for TarFile.add'''

        return (filename[-4:] == '.pyc' or os.path.basename(filename)[0] == '.')

    def __compress(self):
        '''Create the compressed output.

        Should not be used directly, use Compressed.data or Compressed.embeded_run instead
        '''

        output = StringIO()
        tar = TarFile.open(fileobj = output, mode = 'w:bz2')

        for path, modname in self._paths:
            tar.add(path, arcname = modname, exclude = self.exclude_file)

        data = '\n'.join(self._dependencies)
        depinfo = TarInfo(name = 'DEPENDENCIES')
        depinfo.size = len(data)
        depinfo.mtime = time.time()
        depinfo.mode = stat.S_IRUSR | stat.S_IWUSR | stat.S_IRGRP | stat.S_IROTH
        tar.addfile(depinfo, fileobj = StringIO(data))

        tar.close()
        self._data = output.getvalue()

    @property
    def base64(self):
        '''Get the data in base64 format'''

        return b64encode(self.data)

    def split_lines(self, line_width = 80):
        '''Split the base64 encoded output to separate lines'''

        tar = self.base64

        res = []
        while tar:
            line = tar[:line_width]
            tar = tar[line_width:]
            res.append(line)

        data = '\n'.join(res)

    def embeded_run(self, output, module, function, name):
        '''Get a python self-extracting script to run

        Arguments:
        output -- File object destination
        module -- Main module name
        function -- Function to call from main module
        name -- Name of the archive
        '''

        checksum = self.checksum
        separator = '# %s' % checksum
        print >>output, '''#!/usr/bin/env python

def check_decompress():
    try:
        from cStringIO import StringIO

    except ImportError:
        from StringIO import StringIO

    import os
    import os.path
    import tarfile
    import sys
    import base64

    try:
        import %(module)s
        return

    except ImportError:
        pass

    outdir = '.%(name)s-%(checksum)s'
    outdir = os.path.join(os.path.dirname(__file__), outdir)
    separator = '%(separator)s'
    if not os.path.isdir(outdir):
        if os.path.exists(outdir):
            raise Exception("%%s exists but is not a directory" %% outdir)

        myself = open(__file__).read()
        pos = myself.rfind(separator)
        if pos == -1:
            raise Exception("Separator not found in file")

        data = myself[pos+len(separator):].strip()
        data = base64.b64decode(data)

        os.mkdir(outdir)
        inobj = StringIO(data)
        tarfile.open(fileobj = inobj, mode = 'r:*').extractall(path = outdir)

        # Check dependencies
        depsname = os.path.join(outdir, 'DEPENDENCIES')
        fdeps = open(depsname)
        deps = fdeps.read().split('\\n')
        fdeps.close()

        for dep in deps:
            try:
                __import__(dep)

            except ImportError:
                print >>sys.stderr, "Warning: missing dependency %%s" %% dep

    sys.path.insert(0, outdir)

check_decompress()

from %(module)s import %(function)s
%(function)s()
''' % locals()

        print >>output, separator, self.base64

def main():
    compress = Compressed()
    compress.add_module('containment')

    fout = open('testout.py', 'w')
    compress.embeded_run(fout, 'containment', 'main', 'containment')
    fout.close()
    os.system('chmod 755 testout.py')

    fout = open('testout-%s.tar.bz2' % compress.checksum, 'w')
    fout.write(compress.data)
    fout.close()

if __name__ == '__main__':
    main()