#!/usr/bin/python
# python script that compiles files and directories to python bytecode,
# possibly with optimizations

# TODO:
#   Distinguish between files and directories
#   Accept the optimization flag
#   Enforce command option rules

import py_compile
import compileall
import sys
import os
import re
import getopt
import traceback

usage = r"""
pycc : Python byteCode Compiler (pronounced "pie-sees")
(C) Andrew Kesterson 2006, <andrew@aklabs.net>
usage: pycc <options> src_files ...
options:
    -L dir  :   Add the following directory to the sys.path; similar to
                the -L directive in gcc
    -l lib  :   Specifically import the given module before compiling 
                any modules (successive -l directives are processed in 
                the order received); similar to -l directive in gcc. The
                argument should be in standard python 'import' notation. 
                "from xxx import y' should be enclosed in quotes. Otherwise 
                provide 'xxx' alone for the module name to import
    -O      :   Enable the python compiler's optimizations; default output
                will have a .pyo extension instead of a .pyc. NOTE: FOR
                SOME REASON, this option doesn't work on single files.
                Only directories. This is a limitation of the python
                compiler libraries, not this program.
    -o      :   When compiling a single file, specify the output filename. 
                When  compiling a directory or multiple files, this directive 
                has no effect.
    -r      :   Recurse beyond the first level of directories given on the 
                command line, so that for example passing src/dir on the 
                command line would compile src/dir/src1.py and also 
                src/dir/subdir/src2.py.
    src_files   This is a list of single files and directories. Directories 
                are not recursed into (beyond the first level) unless the 
                -r option is specified.
    -E     :    Stop compilation of all source files when one source file 
                has an error (otherwise, an error message is printed to 
                stderr and compilation continues to the next file). Error
                messages are printed regardless of this flag.
    -d n    :   Tells the compiler to descend up to (at maximum) n levels of 
                directories when running recursively
    -F      :   Force the generation of code, even when the .pyc/.pyo 
                timestamps are up to date
    -v      :   Be verbose (errors are always printed regardless)
    -h      :   display this help
"""

shortopts = "L:l:Oo:rd:EvhF"

class ArgChecker:
    def printErr(self, msg):
        if ( not msg.endswith("\n") ):
            msg += "\n"
        sys.__stdout__.write(msg)

    def importPaths (self, importStr):
        fromRE = re.compile(r"""^from\s+([\w.]+)\s+import\s+(\w+).*$""")
        importRE = re.compile(r"""^import\s+([.\w]+).*$""")
        fromparts = fromRE.findall(importStr)
        importparts = importRE.findall(importStr)
        try:
            if ( (len(fromparts) != 1 or len(fromparts[0]) != 2 ) and len(importparts) == 0):
                self.printErr("Bad option: -l %s : invalid syntax : not enough arguments!" % (importStr))
                self.printErr("%s : %s" % (fromparts, importparts))
                return False
            elif ( len(fromparts) == 1 and len(fromparts[0]) == 2 ):
                __import__(fromparts[0][1], globals(), None, fromparts[0][0].split("."))
                return True
            elif ( len(importparts) == 1):
                __import__(importparts[0][0], globals(), None, None)
                return True
            return False
        except Exception, e:
            self.printErr("Import error:%s: %s" % (importStr, str(e))) 
            return False

    def checkArgs (self, argc, argv):
        global shortopts
        global usage
        compiler = PYCompiler()
        opts, args = getopt.getopt(argv, shortopts)
#        self.printErr("%s : %s" % (opts, args))
        errors = False
        if ( len(args) == 0 ):
            self.printErr(usage)
            return None
        for file in args:
            if ( os.path.exists(file) ):
                if ( os.path.isfile(file) ):
                    compiler.queueFile(file)
                elif ( os.path.isdir(file) ):
                    compiler.queueDir(file)
                else:
                    self.printErr("Bad Option: Unable to determine type of file %s" % file)
                    errors = True
            else:
                self.printErr("Bad Option: File or directory does not exist %s" % file)
                errors = True
        for opt in opts:
            if ( opt[0] == "-L" ):
                # add library path
                if ( os.path.exists(opt[1]) and os.path.isdir(opt[1]) ):
                    sys.path.append(opt[1])
                else:
                    self.printErr("Bad option: -L %s : Path is not a directory or is nonexistant" % opt[1])
                    errors = True
            elif ( opt[0] == "-l" ):
                # specifically import the given module before continuing
                if ( not self.importPaths(opt[1]) ): 
                    errors = True
            elif ( opt[0] == "-O" ) :
                compiler.optimize(True)
            elif ( opt[0] == "-o" ) :
                if ( len(compiler.files) == 1 and  len(compiler.dirs) == 0 ):
                    compiler.ofile = opt[1]
                else:
                    if ( compiler.verbose ):
                        self.printErr("Ignoring option -o with multiple files or directories...")
            elif ( opt[0] == "-r" ):
                compiler.recursive(True)
            elif ( opt[0] == "-E" ):
                compiler.stopOnError(True)
            elif ( opt[0] == "-v" ):
                compiler.setVerbose(True)
            elif ( opt[0] == "-F" ):
                compiler.forceGen(True)
            elif ( opt[0] == "-h" ):
                self.printErr("HELP FOUND")
                return False
            elif ( opt[0] == "-d" ):
                try:
                    compiler.setMaxDepth(int(opt[1]))
                except Exception, e:
                    self.printErr("Bad Option: %s : must pass an integer" % opt[0])
            else:
                self.printErr(usage)
                errors = True
        if ( not errors ):
            return compiler
        else:
            #self.printErr("%s : %s" % (opts, args))
            return None

class PYCompiler :
    def __init__ (self):
        self.files = []
        self.dirs = []
        self.opt = False
        self.recurse = False
        self.stopErr = False
        self.verbose = 0
        self.depth = 10
        self.force = 0
        self.ofile = ""

    def setVerbose (self, opt):
        if ( opt ):
            self.verbose = 1
        else:
            self.verbose = 0

    def forceGen (self, opt):
        if ( opt ): 
            self.force = 1
        else:
            self.force = 0

    def setMaxDepth (self, depth):
        self.depth = depth

    def printErr(self, msg):
        sys.__stdout__.write(msg)
        if ( (not msg.endswith("\n")) ):
            sys.__stdout__.write("\n")

    def queueFile (self, fname):
        self.files.append(fname)

    def queueDir (self, dirname):
        self.dirs.append(dirname)

    def printCompileErr (self, e):
        #self.printErr("printErr called with exception %s" % str(e))
        efile = e.file
        etype = e.exc_type_name
        if ( len(e.exc_value) > 1 ) :
            eline = e.exc_value[1][1]
            ecode = e.exc_value[1][3]
        else:
            eline = -1
            ecode = "(no code given)"
        edesc = e.exc_value[0]
        msg = "%s:%d: %s : %s : %s" % (efile, eline, etype, edesc, str(ecode))
        self.printErr(msg)

    def compile (self):
        for fname in self.files:
            ofile = ""
            if ( len(self.files) == 1 and len(self.dirs) == 0 and self.ofile) :
                ofile = self.ofile
            else:
                ofile = fname + "c"
            if ( not self.__compile_file(fname, ofile) ):
                if ( self.stopErr or ( len(self.files) + len(self.dirs)) == 1):
                    return False
        for dirname in self.dirs :
            if ( not os.path.walk(dirname, self.__compile_dir, None) ):
                if ( self.stopErr ):
                    return False
        return True

    def __compile_file (self, fname, ofile = None):
        try:
            if ( self.force == 1 and len(self.files) == 1 and len(self.dirs) == 0):
                if ( os.path.exists(ofile) and os.path.isfile(ofile) ):
                    try:
                        os.remove(ofile)
                    except OSError, e:
                        if ( self.verbose ):
                            self.printErr("%s : Failed to remove original - ignoring ... " % fname)
                else:
                    if ( self.verbose ):
                        self.printErr("%s : Original output does not exist - ignoring ... " % fname)
            #self.printErr("%s -> %s ..." % (fname, ofile))
            py_compile.compile(fname, ofile, None, True)
            if ( (fname in self.files) and self.verbose ):
                # don't print this if we're being called from __compile_dir
                self.printErr("%s -> %s ... OK" % (fname, ofile))
            return True
        except py_compile.PyCompileError, e:
            self.printCompileErr(e)
            if ( fname in self.files ):
                self.printErr("%s -> %s ... FAILED" % (fname, ofile))
            return False

    def __compile_dir (self, arg, dirname, fnames):
        try:
            for fname in fnames:
                if ( not (fname.endswith("py") or fname.endswith("PY")) ):
                    continue
                tocompile = os.path.join(dirname, fname)
                if ( os.path.isdir(tocompile) ):
                    os.path.walk(dirname, self.__compile_dir, None)
                outfile = tocompile+"c"
                if ( not self.__compile_file(tocompile, outfile) ):
                    self.printErr("%s -> %s ... FAILED" % (tocompile, outfile))
                    if ( self.stopErr ):
                        return False
                    else:
                        continue
                else:
                    if ( self.verbose ):
                        self.printErr("%s -> %s ... OK" % (tocompile, outfile))
            return True
        except py_compile.PyCompileError, e:
            self.printCompileErr(e)
            return False

    def optimize (self, opt):
        self.opt = opt

    def recursive (self, opt):
        self.recurse = opt

    def stopOnError (self, opt):
        self.stopErr = opt

def main (argc, argv):
    checker = ArgChecker()
    compiler = checker.checkArgs(argc, argv)
    if ( not compiler ):
        sys.exit(1)
    else:
        if ( not compiler.compile() ):
            sys.exit(2)
        else:
            sys.exit(0)

if ( __name__ == "__main__" ):
    main(len(sys.argv)-1, sys.argv[1:])
