__all__ = ['signature', 'defop', 'overloaded', 'overload']

def signature(*args, **kw):
    """Decorator to set a function's __signature__

    In Py3K (or sooner), this would happen automatically done from argument
    annotations
    """
    def decorate(f):
        f.__signature__ = args, kw
        return f
    return decorate


def defop(gf):
    """Decorator to register a method with a generic function

    In Py3K, this would be replaced by a 'defop' statement.  (Maybe in 2.x if
    a ``__future__`` is used.)
    """
    def decorate(f):
        overload(gf, f)
        return f
    return decorate


class overloaded(object):
    """A simple multi-dispatch implementation, w/caching"""

    def __init__(self, default=None):
        self.registry = {}
        if default:
            self.__name__ = getattr(default,'__name__',None)
            self.__signature__ = getattr(default,'__signature__',None)
            self[None] = default

    def __call__(self, *args):
        """Call the overloaded function."""
        types = tuple(map(type, args))
        func = self.cache.get(types)
        if func is None:
            self.cache[types] = func = self[types]
        return func(*args)

    def __setitem__(self, signature, func):
        self.registry[signature] = func
        self.cache = self.registry.copy()

    def __getitem__(self, types):
        registry = self.registry
        if types in registry:
            return registry[types]
        best = []
        for sig in registry:
            if not implies(types, sig):
                continue # skip inapplicable signatures
            to_remove = []
            for other in best:
                if implies(other, sig):
                    if not implies(sig,other):
                        break   # other dominates us, so don't include
                    # else we are overlapping, so add both
                elif implies(sig,other):
                    # we dominate other, so it has to be removed
                    to_remove.append(other)
                # else we are disjoint, so add both
            else:
                best.append(sig)
            for other in to_remove:
                best.remove(other)
        if len(best)==1:
            return registry[best[0]]

        # Perhaps these multiple candidates all have the
        # same implementation?
        funcs = set(self.registry[sig] for sig in best)
        if len(funcs) == 1:
            return funcs.pop()

        raise TypeError("ambiguous methods; types=%r; candidates=%r" %
            (types, best))


@overloaded
def overload(gf, f):
    """Add a method to an overloaded function"""
    signature,kw = getattr(f,'__signature__',(None,None))
    gf[signature] = f


@overloaded
def implies(types, sig):
    """Do `types` imply `sig`?"""
    return False

# Manually bootstrap tuple/None, anything/None, and type/type comparisons; this
# has to be done before the first call of any overloaded functions, including
# overload() itself, which is called by @defop.  So these bootstraps have to
# happen before we can use @defop.  (Technically, the type/type one doesn't
# *need* to happen until later, but this is a convenient place to put it.)
#
implies[object, type(None)] = lambda s,t: True
implies[tuple,  type(None)] = lambda s,t: True
implies[type,   type]       = issubclass

# Now that that's taken care of, we can use @defop to define everything else

# defop implies(s1:tuple, s2:tuple)
#
@defop(implies)
@signature(tuple, tuple)
def tuple_implies(sig, types):
    if len(sig)!=len(types):
        return False
    for s,t in zip(sig,types):
        if not implies(s,t):
            return False
    return True


# Now we can implement a kind of poor-man's typeclass, wherein we can treat
# a 1-argument generic function as if it were a type.  This implementation is
# pretty crude, as a complete typeclass system should allow more sophisticated
# ways to specify the signature(s).  But that will be left as an exercise for
# the reader at this point.  ;)  Note that with appropriate overloads of
# ``overload()`` and ``implies()``, you can totally define your own framework
# for how overloads are chosen, so e.g. RuleDispatch could just provide
# overloads for these operations that work with its own generic function
# implementation and provide some extra decorators and type annotation objects
# to plug into the base system.


# defop implies(g1:type, g2:overloaded)
#
@defop(implies)
@signature(type, overloaded)
def type_implies_gf(t,gf):
    """Here we implement the equivalent of issubclass(cls,genericfunction)"""
    try:
        return gf[t,] is not None
    except TypeError:
        return False


# defop implies(g1:overloaded, g2:overloaded)
#
@defop(implies)
@signature(overloaded, overloaded)
def gf_implies_gf(g1,g2):
    """Here we implement the equivalent of issubclass(g1,g2)"""
    for s in g1.registry:
        try:
            if g2[s] is None:
                return False
        except TypeError:
            return False
    else:
        return True # g1 implies g2 if g2 has methods for all of g1's sigs


# Doctest suite here, picked up by "setup.py test"

def additional_tests():
    import doctest
    return doctest.DocFileSuite(
        'overloading.txt', optionflags=doctest.ELLIPSIS, package=__name__,
    )