#!/usr/bin/env python
#
# SQUEEZE
#
# squeeze a python program 
#
# installation:
# - use this script as is, or squeeze it using the following command:
# 
# python squeezeTool.py -1su -o squeeze -b squeezeTool squeezeTool.py
#
# notes:
# - this is pretty messy.  make sure to test everything carefully
#   if you change anything
#
# - the name "squeeze" is taken from an ABC800 utility which did
#   about the same thing with Basic II bytecodes.
#
# history:
# 1.0   97-04-22 fl     Created
# 1.1   97-05-25 fl     Added base64 embedding option (-1)
#       97-05-25 fl     Check for broken package file
# 1.2   97-05-26 fl     Support uncompressed packages (-u)
# 1.3   97-05-27 fl     Check byte code magic, eliminated StringIO, etc.
# 1.4   97-06-04 fl     Removed last bits of white space, removed try/except
# 1.5   97-06-17 fl     Added squeeze archive capabilities (-x)
# 1.6   98-05-04 fl     Minor fixes in preparation for public source release
#
# reviews:
#       "Fredrik Lundh is a friggin genius"
#       -- Aaron Watters, author of 'Internet Programming with Python'
#
#       "I agree ... this is a friggin Good Thing"
#       -- Paul Everitt, Digital Creations
#
# Copyright (c) 1997 by Fredrik Lundh.
# Copyright (c) 1997-1998 by Secret Labs AB
#
# info@pythonware.com
# http://www.pythonware.com
#
# --------------------------------------------------------------------
# Permission to use, copy, modify, and distribute this software and
# its associated documentation for any purpose and without fee is
# hereby granted.  This software is provided as is.
# --------------------------------------------------------------------

VERSION = "1.6/98-05-04"
MAGIC   = "[PANDASQUEEZE]"

import base64, imp, marshal, os, string, sys, md5

# --------------------------------------------------------------------
# usage

def usage():
	print
	print "SQUEEZE", VERSION, "(c) 1997-1998 by Secret Labs AB"
	print """\
Convert a Python application to a compressed module package.

Usage: squeeze [-1ux] -o app [-b start] modules... [-d files...]

This utility creates a compressed package file named "app.pyz", which
contains the given module files.  It also creates a bootstrap script
named "app.py", which loads the package and imports the given "start"
module to get things going.  Example:

        squeeze -o app -b appMain app*.py

The -1 option tells squeeze to put the package file inside the boot-
strap script using base64 encoding.  The result is a single text file
containing the full application.

The -u option disables compression.  Otherwise, the package will be
compressed using zlib, and the user needs zlib to run the resulting
application.

The -d option can be used to put additional files in the package file.
You can access these files via "__main__.open(filename)" (returns a
StringIO file object).

The -x option can be used with -d to create a self-extracting archive,
instead of a package.  When the resulting script is executed, the
data files are extracted.  Omit the -b option in this case.
"""
	sys.exit(1)


# --------------------------------------------------------------------
# squeezer -- collect squeezed modules

class Squeezer:

	def __init__(self):

		self.rawbytes = self.bytes = 0
		self.modules = {}

	def addmodule(self, file):

		if file[-1] == "c":
			file = file[:-1]

		m = os.path.splitext(os.path.split(file)[1])[0]

		# read sourcefile
		f = open(file)
		codestring = f.read()
		f.close()

		# dump to file
		self.modules[m] = compile(codestring, m, "exec")

	def adddata(self, file):

		self.modules["+"+file] = open(file, "rb").read()

	def getarchive(self):

		# marshal our module dictionary
		data = marshal.dumps(self.modules)
		self.rawbytes = len(data)

		# return (compressed) dictionary
		data = zlib.compress(data, 9)
		self.bytes = len(data)

		return data

	def getstatus(self):
		return self.bytes, self.rawbytes


# --------------------------------------------------------------------
# loader (used in bootstrap code)

loader = """
import ihooks

PYZ_MODULE = 64

class Loader(ihooks.ModuleLoader):

    def __init__(self, modules):
        self.__modules = modules
        return ihooks.ModuleLoader.__init__(self)

    def find_module(self, name, path = None):
        try:
            self.__modules[name]
            return None, None, (None, None, PYZ_MODULE)
        except KeyError:
            return ihooks.ModuleLoader.find_module(self, name, path)

    def load_module(self, name, stuff):
        file, filename, (suff, mode, type) = stuff
        if type != PYZ_MODULE:
            return ihooks.ModuleLoader.load_module(self, name, stuff)
        #print "PYZ:", "import", name
        basename = name.split('.')[-1]
        code = self.__modules[basename]
        del self.__modules[basename] # no need to keep this one around
        m = self.hooks.add_module(name)
        m.__file__ = filename
        exec code in m.__dict__
        return m

def boot(name, fp, size, offset = 0):

    global data

    try:
        import %(modules)s
    except ImportError:
        #print "PYZ:", "failed to load marshal and zlib libraries"
        return # cannot boot from PYZ file
    #print "PYZ:", "boot from", name+".PYZ"

    # load archive and install import hook
    if offset:
        data = fp[offset:]
    else:
        data = fp.read(size)
        fp.close()

    if len(data) != size:
        raise IOError, "package is truncated"

    data = marshal.loads(%(data)s)

    ihooks.install(ihooks.ModuleImporter(Loader(data)))
"""

loaderopen = """

def open(name):
    import StringIO
    try:
        return StringIO.StringIO(data["+"+name])
    except KeyError:
        raise IOError, (0, "no such file")
"""

loaderexplode = """

def explode():
    for k, v in data.items():
        if k[0] == "+":
            try:
                open(k[1:], "wb").write(v)
                print k[1:], "extracted ok"
            except IOError, v:
                print k[1:], "failed:", "IOError", v

"""

def getloader(data, package):

	s = loader

	if data:
		if explode:
			s = s + loaderexplode
		else:
			s = s + loaderopen

	dict = {
		"modules": "marshal, zlib",
		"data":    "zlib.decompress(data)",
		}

	s = s % dict

	return marshal.dumps(compile(s, "<package>", "exec"))


# --------------------------------------------------------------------
# Main
# --------------------------------------------------------------------

#
# parse options

import getopt, glob, sys
import zlib

embed = 0
explode = 0

def squeeze(app, start, filelist, outputDir):
	localMagic = MAGIC
	data = None

	bootstrap = os.path.join(outputDir, app + ".py")
        archiveBase = app + ".pyz"
	archive   = os.path.join(outputDir, archiveBase)

	archiveid = app

	#
	# avoid overwriting files not generated by squeeze

	try:
		fp = open(bootstrap)
		s = fp.readline()
		string.index(s, MAGIC)
	except IOError:
		pass
	except ValueError:
		print bootstrap, "was not created by squeeze.  You have to manually"
		print "remove the file to proceed."
		sys.exit(1)

	#
	# collect modules

	sq = Squeezer()
	for file in filelist:
		# print 'addmodule:', file
		sq.addmodule(file)

	package = sq.getarchive()
	size = len(package)

	#
	# get loader

	loader = getloader(data, package)

	zbegin, zend = "zlib.decompress(", ")"
	loader = zlib.compress(loader, 9)

	loaderlen = len(loader)

	magic = repr(imp.get_magic())
	version = string.split(sys.version)[0]

	#
	# generate script and package files

	if embed:

		# embedded archive
		data = base64.encodestring(loader + package)

		fp = open(bootstrap, "w")
		fp.write('''\
#%(localMagic)s %(archiveid)s
import ihooks,zlib,base64,marshal
s=base64.decodestring("""
%(data)s""")
exec marshal.loads(%(zbegin)ss[:%(loaderlen)d]%(zend)s)
boot("%(app)s",s,%(size)d,%(loaderlen)d)
exec "import %(start)s"
''' % locals())
		bytes = fp.tell()

	else:

		# separate archive file

		fp = open(archive, "wb")

		fp.write(loader)
		fp.write(package)

		bytes = fp.tell()
		fp.close()
		#
		# create bootstrap code

		fp = open(bootstrap, "w")
		# Note: David Rose adjusted the following to be panda-specific.
		fp.write("""\
#%(localMagic)s %(archiveid)s
import ihooks,zlib,marshal,os,sys

import pandac

def searchPath(filename):
  # Look along pandac.__path__ for the indicated filename.  Returns
  # the located pathname, or None if the filename is not found.
  for dir in pandac.__path__:
    pathname = os.path.join(dir, filename)
    if os.path.exists(pathname):
      return pathname

  return None

# Look for %(archiveBase)s along pandac.__path__.
archiveName = "%(archiveBase)s"
archivePath = searchPath(archiveName)
if archivePath == None:
  raise ImportError, "Could not locate pandac.%%s." %% (archiveName)

f=open(archivePath,"rb")
exec marshal.loads(%(zbegin)sf.read(%(loaderlen)d)%(zend)s)
boot("%(app)s",f,%(size)d)
exec "from %(start)s import *"
#exec "run()"
""" % locals())
		bytes = bytes + fp.tell()
		fp.close()

	#
	# show statistics

	dummy, rawbytes = sq.getstatus()

	print "squeezed", rawbytes, "to", bytes, "bytes",
	print "(%d%%)" % (bytes * 100 / rawbytes)
