#!/usr/bin/python
# $Id: $

'''annotype.py

Usage::

    python -m annotype [options] <script> [script-args ...]

Options:

    --help      Display this message
    --dir=PATH  Set output directory (defaults to "out_annotype")
    --all       Do not discriminate files based on path.  By default,
                files with an absolute path are ignored (e.g., Python
                standard library).

Example::

    python -m annotype --dir=out examples/opengl.py

Runs the pyglet example "opengl.py", writing an annotated source file to
directory "out".
'''

import compiler
import os
import re
import sys
import traceback

sources = {}

class Source(object):
    def __init__(self, path):
        self.path = path

        # Trim path down
        dir = ''
        path, filename = os.path.split(path)
        while len(dir + filename) < 30:
            filename = os.path.join(dir, filename)
            path, dir = os.path.split(path)
            if not dir:
                filename = os.path.join('', filename)
                break
        else:
            filename = os.path.join('...', filename)
        self.abbrv_path = filename

        self.funcs = {}

    def write(self, file):
        for i, line in enumerate(open(self.path)):
            try:
                func = self.funcs[i]
                if func.argtypes:
                    indent = _get_indent(line)
                    for name in func.argnames:
                        file.write('%s# [annotype] %s: %s\n' % (
                            indent, name, ', '.join(func.argtypes[name])))
            except KeyError:
                pass
            file.write(line)

class Func(object):
    def __init__(self, code):
        self.line = code.co_firstlineno
        self.name = code.co_name
        self.argtypes = {}
        self.argnames = [] # Same as keys of argtypes, but ordered

_indent_re = re.compile('^(\s*).*')
def _get_indent(line):
    return _indent_re.match(line).groups()[0]

def _trace_frame(frame):
    code = frame.f_code
    path = code.co_filename

    if path.startswith('/') and not trace_all:
        return

    try:
        source = sources[path]
    except KeyError: 
        source = sources[path] = Source(path)

    try:
        func = source.funcs[code.co_firstlineno]
    except KeyError:
        func = source.funcs[code.co_firstlineno] = Func(code)

    for argname in code.co_varnames[:code.co_argcount]:
        argvalue = frame.f_locals[argname]
        try:
            types = func.argtypes[argname]
        except KeyError:
            types = func.argtypes[argname] = set()
            func.argnames.append(argname)
        types.add(type(argvalue).__name__)

def _trace_func(frame, event, arg):
    if event == 'call':
        _trace_frame(frame)

# Don't pull this into a function: exec global name lookups broken when exec
# is called from within a function.
if __name__ == '__main__':
    dir = 'out_annotype'
    trace_all = False

    while len(sys.argv) > 1 and sys.argv[1].startswith('--'):
        option = sys.argv.pop(1)
        if '=' in option:
            k, v = option.split('=')
        else:
            k, v = option, None

        if k == '--dir':
            dir = v
        elif k == '--all':
            trace_all = True
        elif k == '--help':
            print __doc__
            sys.exit(0)

    if len(sys.argv) < 2:
        print __doc__
        sys.exit(0)

    del sys.argv[0]

    code = compiler.compile(''.join(open(sys.argv[0])), sys.argv[0], 'exec')

    sys.setprofile(_trace_func)
    try:
        exec code
        sys.setprofile(None)
    except:
        sys.setprofile(None)
        traceback.print_exc()

    try:
        os.makedirs(dir)
    except OSError:
        pass

    for source in sources.values():
        out = open(os.path.join(dir, source.path.replace('/', '_')), 'w')
        source.write(out)