#!/usr/bin/python2.4
#
# Copyright 2007 Google Inc. All Rights Reserved.

"""Script to help with making portable, executable .pyz zip files.

Here is the pipeline:

   Source file tree

-> [ zip action ] ->

   Zip file with sources in it

-> [ pyz action ] ->

   .pyz file with a __zipmain__.py module at its root, executable with python -z

-> [ unix action ] ->

   .pyz file with shebang line, directly executable on Unix-like systems

On non-Unix platforms, the pyz action produces a .pyz file already executable
by the Python interpreter.
"""


import optparse
import os
import sys
import zipfile


class UsageError(RuntimeError):
    pass


def makezip(source_files, output_zip):
    """Create a new zip file from source files."""
    out = zipfile.ZipFile(output_zip, 'w')
    for filename in source_files:
        out.write(filename)
        print 'Added %s to %s' % (filename, output_zip)
    out.close()


_ZIPMAIN_TEMPLATE = """
import sys
%(import_str)s

sys.exit(%(module)s.main(sys.argv))
"""

def add_zipmain(package, module, output_zip):
    """Turn a zip file into an executable .pyz file by adding a __zipmain__.py
    file.
    """
    out = zipfile.ZipFile(output_zip, 'a') # append mode
    if options.package:
        import_str = 'from %s import %s' % (options.package, options.module)
    else:
        import_str = 'import %s' % options.module

    out.writestr('__zipmain__.py', _ZIPMAIN_TEMPLATE %
        {'import_str': import_str, 'module': options.module})
    out.close()
    print 'Added __zipmain__.py to %s' % options.output_zip
    

def make_unix_executable(python, mask, output_zip):
    """Add a shebang line and chmod the file for direct execution on Unix
    systems.
    """
    shebang = '#!%s -z' % python

    f = open(output_zip, 'rb+')
    contents = f.read()
    f.seek(0)
    f.write(shebang + '\n')
    f.write(contents)
    f.close()
    print 'Prepended %s to %s.' % (shebang, output_zip)

    os.chmod(output_zip, mask)
    print 'chmod %s 0%o' % (output_zip, mask)


def main(options, argv):
    if options.actions:
        options.actions = options.actions.split(',')
    else:
        raise UsageError('Actions required.')

    if not options.output_zip:
        raise UsageError('An output zip filename is required.')

    # Ensure that the actions are executed in order.
    for action in options.actions:
        if action == 'zip':
            if not options.basedir and not argv:
                raise UsageError(
                    'zip action requires a base directory or a list of files.')

            if options.basedir:
                # TODO: do find
                pass
            else:
                source_files = argv
            makezip(source_files, options.output_zip)

        elif action == 'pyz':
            if not options.module:
                raise UsageError('pyz action requires a module name.')
            add_zipmain(options.package, options.module, options.output_zip)

        elif action == 'unix':
            make_unix_executable(
                    options.python, options.mask, options.output_zip)
    return 0


if __name__ == '__main__':
    parser = optparse.OptionParser()

    # For all steps
    parser.add_option(
        '-a', '--actions', dest='actions', type='str',
        help='Comma separated list of actions to execute: zip,pyz,unix')
    parser.add_option(
        '-z', '--zip', dest='output_zip', type='str',
        help='Output zip file')

    # For zip step
    parser.add_option(
        '-d', '--dir', dest='basedir', type='str',
        help='Base directory of the source tree to zip up.')

    # For pyz step
    parser.add_option(
        '-p', '--package', dest='package', type='str', 
        help='Package where the main module lives, e.g. spam.eggs')
    parser.add_option(
        '-m', '--module', dest='module', type='str', 
        help='Module which contains a main() function to execute.')

    # For unix step
    parser.add_option(
        '-y', '--python', dest='python', type='str', default='/usr/bin/python',
        help='Location of the Python interpreter, for putting in the shebang '
        'line')
    parser.add_option(
        '-c', '--chmod', dest='mask', type='int', default=0700,
        help='On Unix systems, chmod the output using this mask.  Make sure '
        'to include a leading 0 for octal.')

    if len(sys.argv) == 1:
        parser.print_help()
        sys.exit(1)

    (options, argv) = parser.parse_args()
    try:
        exit_code = main(options, argv)
    except UsageError, e:
        print >>sys.stderr, 'Usage error: %s' % e
        exit_code = 1
    sys.exit(exit_code)