Pythran-based dummy JIT compiler

jit.py

import pythran
import inspect
import hashlib
import itertools
import imp
import re
import os
import numpy as np

activate_jit = True

def typename(obj):
    # FIXME: only works for some types
    if type(obj) is np.ndarray:
        tmplt = "{dtype}{dim}"
        return tmplt.format(dtype=obj.dtype.name,
                            dim="[]" * obj.ndim)
    return type(obj).__name__

class jit(object):

    def __init__(self, **flags):
        self.flags = flags

    def __call__(self, fun):
        global activate_jit
        if not activate_jit:
            return fun
        # FIXME: gather global dependencies using pythran.analysis.imported_ids
        src = inspect.getsource(fun)
        src = re.sub("@.*?\sdef\s","def ", src)
        print(src)
        fname   = fun.__name__
        m = hashlib.md5()
        m.update(src.encode("utf8"))

        if not os.path.exists("__pythran__"):
            os.makedirs("__pythran__")

        def type_collector(*args, **kwargs):
            arg_types = [typename(arg) for arg in itertools.chain(args, kwargs.values())]
            m.update("".join(arg_types).encode("utf8"))
            func_sig = "{}({})".format(fname, ", ".join(arg_types))
            header = "#pythran export {}\n".format(func_sig)
            print(header)
            # FIXME: implement a cache here

            # module_name = "{}__{}__".format(fname, "_".join(arg_types)) + m.hexdigest()
            module_name = "fname_" + m.hexdigest()
            # FIXME: force output in tmp dir
            module_path = pythran.compile_pythrancode(module_name, header + src, 
                                                      output_file='__pythran__/' + module_name + '.so',
                                                      **self.flags)
            module = imp.load_dynamic(module_name, module_path)
            return getattr(module, fun.__name__)(*args, **kwargs)

        return type_collector

@jit()
def apb(a, b):
    return a + b

if __name__ == '__main__':
    a = np.array([[1,2,3,4],[5,6,7,8]])
    print(apb(a, 123))
    print(apb(1, 4))
    print(apb(1, 3.123))