overload.py

#-*-coding: utf-8;-*-
"""
Overloaded function for python
"""

import enum

class _cfs(enum.Enum):
    OPTIMIZED = 0x001
    NEWLOCALS = 0x002
    VARARGS = 0x004
    VARKEYWORDS = 0x008
    NESTED = 0x010
    GENERATOR = 0x020
    NOFREE = 0x040
    COROUTINE = 0x080
    ITERABLE_COROUTINE = 0x100

class vararg:
    pass

class kwvararg(vararg):
    pass

class _func:
    def __init__(self, func):
        if func.__defaults__ is not None or func.__kwdefaults__ is not None:
            raise TypeError("Overload function can't set default arguments.")
        self.flags = _get_flags(func)
        self.args = _gawtbf(func, self.flags)
        self.func = func
        self.__name__ = func.__name__
    def match(self, *args, **kwargs):
        kwa = kwargs.copy()
        kwo = self.args['keywordonlys'].copy()
        pas = self.args['positionals'].copy()
        last = float('infinity')
        for k in kwargs.keys():
            if k in kwo.keys():
                if not isinstance(kwa[k], kwo[k]):
                    return False
                del kwo[k], kwa[k]
                continue
            if k in pas.keys():
                if not isinstance(kwa[k], pas[k][1]):
                    return False
                del pas[k], kwa[k]
                continue
            if '**' in kwo.keys():
                del kwa[k]
                continue
            return False
        del pas, kwo
        pas = dict(self.args['positionals'].values())
        if len(pas) > len(args):
            return False
        for i in range(len(args)):
            try:
                if not isinstance(args[i], pas[i]):
                    return False
            except IndexError:
                if not -1 in pas.keys():
                    return False
        return True
    def __call__(self, *args, **kwargs):
        return self.func(*args, **kwargs)

class _overloaded:
    def __init__(self, func):
        self.__name__ = func.__name__
        self.funcs = [_func(func)]
    def __call__(self, *a, **kwa):
        for f in self.funcs:
            if not f.match(*a, **kwa):
                continue
            return f(*a, **kwa)
        raise TypeError(
            "Overload function {} matched to the argument type was not found.".format(self.__name__))
    def append(self, func):
        self.funcs.append(_func(func))

def overload(func):
    olfo = globals().get(func.__name__, None)
    if olfo is None:
        return _overloaded(func)
    olfo.append(func)
    return olfo

def _get_flags(func):
    flgs = bin(func.__code__.co_flags)[:1:-1]
    ret = []
    try:
        if flgs[0] == '1': ret.append(_cfs.OPTIMIZED)
        if flgs[1] == '1': ret.append(_cfs.NEWLOCALS)
        if flgs[2] == '1': ret.append(_cfs.VARARGS)
        if flgs[3] == '1': ret.append(_cfs.VARKEYWORDS)
        if flgs[4] == '1': ret.append(_cfs.NESTED)
        if flgs[5] == '1': ret.append(_cfs.GENERATOR)
        if flgs[6] == '1': ret.append(_cfs.NOFREE)
        if flgs[7] == '1': ret.append(_cfs.COROUTINE)
        if flgs[8] == '1': ret.append(_cfs.ITERABLE_COROUTINE)
    except IndexError:
        pass
    return ret

def _gawtbf(func, flags = None):
    if flags is None:
        flags = _get_flags(func)
    ann = func.__annotations__
    co = func.__code__
    pargs = co.co_varnames[:co.co_argcount]
    kwargs = co.co_varnames[co.co_argcount:co.co_argcount+co.co_kwonlyargcount]
    pas = {}
    kwas = {}
    for p in range(len(pargs)):
        a = pargs[p]
        pas[a] = p, ann.get(a, object)
    if _cfs.VARARGS in flags:
        pas["*"] = -1, vararg
    for kwa in kwargs:
        kwas[kwa] = ann.get(kwa, object)
    if _cfs.VARKEYWORDS in flags:
        kwas["**"] = kwvararg
    return {
        "positionals": pas,
        "keywordonlys": kwas}