ice.py

#! /usr/bin/env python3
"""An alternative Python freeze implementation.

The two primary motivations for this implementation of freeze (when
compared to the standard freeze distributed with Python) are

1) Correctly works with Python3.

2) Includes the entire standard library (rather than searching
for specific modules).

The primary reason for this is that the applications I wish to
target allow some form of plugins, so it is not possible to
know the entire set of required modules ahead of time.

Obviously there is a space trade-off and this may not be the best
solution in other use-cases. The standard library, is relatively
large, so if you don't use much of it, and don't need plugins, then
this is probably not for you (try cxFreeze).


Additional motivations:

For all I know cxFreeze, or some other tool out there already supports
this use case (or could be shoe-horned in to supporting it). My other
motivations for creating this are as a learning experience to ensure
that I fully understand the workings of Python frozen modules. Finally
this provides a simple code base for additional experimentation.


Frozen Library:

This module does two main things:

1) Create `frozen libraries`

2) Create an application 'main.c' from a list of frozen libraries.

A frozen library is a group of individual modules frozen in a single
C file. The C file basically contains a C array representation of
the compiled (and optimised) byte code for the module. Additionally
the frozen library has a <library_name>_struct.h and
<library_name>_extern.h file that contains structure elements,
and extern definitions of the library. These header files are
used when creating an app, but can be included in a C file directly
if you know what you are doing.


Usage:

The primary interfaces that will be used are:

  create_stdlib: Creates a frozen version of the standard library.
  create_lib: Create a frozen version of a an existing library.
  create_app: Create a main.c suitable for intialising a list of frozen libraries.


Static Python:

My end goal when using a freeze application is to ensure that I can
distribute a single binary with minimal external dependencies. To
achieve this it is ideal if the targetted Python does not need to use
dynamic extensions. This can be achieved by statically compiling
Python. The exact mechanics for that are beyond the scope fo this
documentation, however the `xyz` packaging tool can be used to
compile an appropriate Python: https://github.com/BreakawayConsulting/xyz
Or if (like me) you prefer full control over what you are doing
see the specific rules and Setup.dist files:

https://github.com/BreakawayConsulting/xyz/blob/master/rules/python.py
https://github.com/BreakawayConsulting/xyz/blob/master/rules/pySetup.dist.darwin

Drawbacks:

This current version does not really have any explicit support for
libraries outside of the standard library. It wouldn't be hard to add.

The command line interface is relatively limited right now. You will
probably get more out of it using it as a library, than as a raw tool.

Currently this tool is just about creating the raw C files. It doen't
create any Makefile or similar for compiling the frozen libraries or
application. This is intentional as my primary use cases involved
systems and application embedded in existing programs with their
own build system.

This has only been tested with my use-cases so your milage may vary!

"""
import marshal
import os
import re
import sys
import tokenize

STDLIB_EXCLUDE_LIST = ['test', 'tkinter', 'turtledemo',
                       'tkinter', 'lib2to3', 'idlelib',
                       'distutils']

SUFFIX = '.py'

HEADER = """/* Generated by freeze.py */
#include <Python.h>

"""

MAIN_FN = """
#include <locale.h>

int
main(int argc, char *argv[])
{
    wchar_t **argv_copy = (wchar_t **)PyMem_Malloc(sizeof(wchar_t*)*(argc + 1));
    /* We need a second copies, as Python might modify the first one. */
    wchar_t **argv_copy2 = (wchar_t **)PyMem_Malloc(sizeof(wchar_t*)*(argc + 1));
    int i, n, res = 0;
    char *oldloc;

    if (!argv_copy || !argv_copy2)
    {
        fprintf(stderr, "out of memory\\n");
        return 1;
    }

    oldloc = strdup(setlocale(LC_ALL, NULL));
    setlocale(LC_ALL, "");
    for (i = 0; i < argc; i++)
    {
        argv_copy[i] = _Py_char2wchar(argv[i], NULL);
        if (!argv_copy[i])
        {
            free(oldloc);
            fprintf(stderr, "Fatal Python error: "
                            "unable to decode the command line argument #%i\\n",
                            i + 1);
            return 1;
        }
        argv_copy2[i] = argv_copy[i];
    }
    argv_copy2[argc] = argv_copy[argc] = NULL;

    setlocale(LC_ALL, oldloc);
    free(oldloc);

    Py_FrozenFlag = 1; /* suppress errors from getpath.c */

    Py_Initialize();
    PySys_SetArgvEx(argc, argv_copy, 0);

    n = PyImport_ImportFrozenModule("__main__");

    if (n < 0)
    {
        PyErr_Print();
        res = 1;
    }

    Py_Finalize();
    for (i = 0; i < argc; i++) {
        PyMem_Free(argv_copy2[i]);
    }
    PyMem_Free(argv_copy);
    PyMem_Free(argv_copy2);

    return res;
}

"""

ARRAY_HEADER = """
static struct _frozen _PyImport_FrozenModules[] = {
"""

ARRAY_TRAILER = """\
    {0, 0, 0} /* sentinel */
};

struct _frozen *PyImport_FrozenModules = _PyImport_FrozenModules;
"""


def dict_inverse(dct, exact=False):
    """Given an input dictionary (`dct`), create a new dictionary
    where the keys are indexed by the values. In the original
    dictionary multiple keys may reference the same value, so the
    values in the new dictionary is a list of keys. The order of keys
    in the list is undefined.

    Example:

    > dict_inverse({1: 'a', 2: 'a', 3: 'c'})
    { 'a': [1, 2], 'c': [3]

    If `dct` has an exact inverse mapping `exact` can be passed as
    True. In this case, the values will be just the original key (not
    a list).

    Example:

    > dict_inverse({1: 'a', 2: 'b', 3: 'c'}, exact=True)
    { 'a': 1, 'b': 2, 'c': 3}

    Note: No checking is done when exact is True, so in the case
    where there are multiple keys mapping the same value it is
    undefined as to which key the value will map to.

    From: https://github.com/bennoleslie/pyutil.git

    """
    if exact:
        return {value: key for key, value in dct.items()}

    r = {}
    for key, value in dct.items():
        r.setdefault(value, []).append(key)
    return r


def write_byte_code(outf, var_name, byte_code):
    """Write out `byte_code` as variable `var_name` to a file-like object `outf`."""

    outf.write('unsigned char {}[] = {{'.format(var_name))
    for i in range(0, len(byte_code), 16):
        outf.write('\n    ')
        for c in bytes(byte_code[i:i + 16]):
            outf.write('%d,' % c)
    outf.write('\n};\n')


def create_frozen_lib(name, mods, aliases={}):
    """Create a new *frozen library* called `name`.

    `mods` is a dictionary of module descriptions indexed by module
    name.  A module description is a triple of
      (filename, short_filename, is_pkg).

    This is the same format as returned by `find_modules`.

    A frozen library can contain a number of aliases. Specifically
    that means that a given Python module may be importable under its
    normal name, plus any number of aliases names. The alias is a
    dictionary indexed by the alias name, with the value being the
    name of the underlying module.

    """
    extern_fmt = 'extern unsigned char {}[];\n'
    struct_fmt = '    {{"{}", {}, {}}},\n'
    c_filename = '{}.c'.format(name)
    h_struct_filename = '{}_struct.h'.format(name)
    h_extern_filename = '{}_extern.h'.format(name)

    # Although easier for the user to specify {alias: module}
    # the code is simpler if we have {module: [aliases]}
    aliases = dict_inverse(aliases)

    with open(h_struct_filename, 'w') as h_struct, \
            open(h_extern_filename, 'w') as h_ext, \
            open(c_filename, 'w') as c_out:

        for mod_name in sorted(mods.keys()):
            (filename, short_filename, is_pkg) = mods[mod_name]
            with tokenize.open(filename) as f:
                code = compile(f.read(), short_filename, 'exec', optimize=2)
            var_name = "M_" + "__".join(mod_name.split("."))
            raw_code = marshal.dumps(code)
            size = len(raw_code)
            # Packages are indicated by negative size; this is part of
            # the Pythohn frozen library internals.
            if is_pkg:
                size = -size

            write_byte_code(c_out, var_name, raw_code)

            h_ext.write(extern_fmt.format(var_name))

            h_struct.write(struct_fmt.format(mod_name, var_name, size))

            # Write out all the aliases
            for alias in aliases.get(mod_name, []):
                h_struct.write(struct_fmt.format(alias, var_name, size))


def find_modules(libdir, excluded=[]):
    """Find all the Python modules (and packages) in a given directory `libdir`.

    Note: The 'libdir' is generally something you would otherwise place in
    the Python path. It must not be a package itself. (i.e.: it should not contain
    an __init__.py file.)

    Return a dictionary of (full_filename, short_filename, is_pkg) tuples
    indexed by fully qualified module name.

    full_filename is a complete path to the module.
    short_filename is the path with the `libdir` prefix removed.
    is_pkg is True if the module is a package.

    This only tries to find .py modules.

    Any module that prefix matches a module listed in 'excluded' will
    be ignored.

    """
    # Ensure libdir endswith a trailing slash
    if not libdir.endswith(os.path.sep):
        libdir += os.path.sep

    modules = {}

    if not os.path.exists(libdir) or not os.path.isdir(libdir):
        raise Exception("libdir {} doesn't exist.".format(libdir))

    for root, dirs, fns in os.walk(libdir):
        if "__pycache__" in dirs:  # pycache directories are ignored
            dirs.remove("__pycache__")

        for d in dirs[:]:
            # If a directory doesn't look like a package, ignore it.
            if not os.path.exists(os.path.join(root, d, '__init__' + SUFFIX)):
                dirs.remove(d)

        pkg_name = root[len(libdir):].replace(os.path.sep, '.')
        for fn in [fn for fn in fns if fn.endswith(SUFFIX)]:
            filename = os.path.join(root, fn)
            short_filename = filename[len(libdir):]
            mod_name = fn[:-len(SUFFIX)]
            is_pkg = mod_name == "__init__"
            if is_pkg and pkg_name == '':
                raise Exception("The libdir can't include an __init__.py file.")

            if pkg_name != '':
                mod_name = pkg_name if is_pkg else pkg_name + '.' + mod_name

            exclude = False
            for e in excluded:
                if mod_name.startswith(e):
                    exclude = True
                    break

            if not exclude:
                modules[mod_name] = (filename, short_filename, is_pkg)

    return modules


def merge_dict(dict1, dict2):
    return dict(list(dict1.items()) + list(dict2.items()))


def create_stdlib(excluded=STDLIB_EXCLUDE_LIST):
    """Create a frozen version of the standard library.

    The location of the standard library is determined automatically
    through introspection.

    Specific modules can be excluded with the frozen library. `excluded`
    is a list of modules to exclude (via a simple prefix match).

    """
    version = sys.version[:3]
    libdir = os.path.join(sys.prefix, 'lib', 'python{}'.format(version))
    elibdir = os.path.join(sys.exec_prefix, 'lib', 'python{}'.format(version), 'lib-dynload')
    create_frozen_lib('stdlib',
                      merge_dict(find_modules(libdir, excluded=excluded),
                                 find_modules(elibdir, excluded=excluded)),
                      {'_frozen_importlib': 'importlib._bootstrap'})


def create_lib(name, path, main=None, excluded=[]):
    """Create a new frozen library with modules from the specified
    path. `main` can be set to the name of a module, which will be
    aliases as '__main__' in the frozen library.

    NOTE: The main parameter may be deprecated in the future in favour
    of passing it directly to create_app.

    """
    mods = find_modules(path, excluded=excluded)
    aliases = {}
    if main is not None:
        aliases['__main__'] = main
    create_frozen_lib(name, mods, aliases)


def create_app(libs, filename='main.c'):
    """Create an application from a list of frozen libraries.

    One (and only one) of the frozen libraries should have been
    created with the 'main' argument.

    The application is output in file `filename` (which defaults to
    main.c

    NOTE: In the future create_app may directly specify the module
    to use as __main__, rather than relying on a create_lib.

    """
    with open(filename, 'w') as outf:
        outf.write(HEADER)
        outf.write(MAIN_FN)

        for l in libs:
            with open('{}_extern.h'.format(l)) as f:
                outf.write(f.read())

        outf.write(ARRAY_HEADER)
        for l in libs:
            with open('{}_struct.h'.format(l)) as f:
                outf.write(f.read())

        outf.write(ARRAY_TRAILER)


def main(argv):
    import argparse

    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest='command')

    stdlib_p = subparsers.add_parser('stdlib', help='Create frozen standard library.')

    lib_p = subparsers.add_parser('lib', help='Create frozen library')
    lib_p.add_argument('name', help='library name')
    lib_p.add_argument('path', help='library path')
    lib_p.add_argument('--main', default=None, help='library path')

    app_p = subparsers.add_parser('app', help='Create app main.c from frozen libraries')
    app_p.add_argument('libs', nargs='+', help='libraries')

    args = parser.parse_args()

    if args.command is None:
        parser.print_help()
        return 1

    if args.command == 'stdlib':
        create_stdlib()
    elif args.command == 'lib':
        create_lib(args.name, args.path, args.main)
    elif args.command == 'app':
        create_app(args.libs)

    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))