python 实现有定义域的函数，可用来做类型测试

# -*- coding: utf-8 -*-
# fcool: a cool tool for functional programming
# Author: William Song
#
# CLasses:
#     BaseFunction -> Type (or Domain), Function
#     BaseFunction: func: function (or number)
#     Function: func,
#               domain: Type, the definition domain of func
# Functions:
#     Interval(a:num, b:num) -> Type
#     restrict(t:Type) -> decorator(f:function -> Function(f, t))
# Constants:
#     TURE, FLASE(Type) represent universal set, empty set
#
# Example:
#    F = Function(lambda x:x, Type(lambda x:x>2))
#    F(3)


#    f = 3     # or lambda x,y: 3
#    g = lambda x,y: 2/x
#    t = Type(lambda x:x<5 and isinstance(x, int)) * TRUE  # define type(domain) and functions on it
#        <=> Type(lambda x, y:x<5 and isinstance(x, int))
#    G = Function(g, t)
#    F = Function(f, t)


#    print(G.glue(F)(3,4), glue(G, F)(3,4))    # glue functions


#    ID = Function(lambda x:x)
#    print(ID.compose(F)(3,4))      # composition


#    @restrict(Interval(1,2))       # restricting decorator
#    def f(x):
#        return x


#    print(f(1))
#    try:
#        print(f(3))
#    except Exception as ex:
#        print(ex)


#    G=Function(lambda x:x)
#    print(G(3))
#    G = G | Interval(1,2)           # restricting method   
#    print(G(3))


import operator
import functools
from types import FunctionType




class Arithmetic(object):


    def copy(self):
        return NotImplemented


    def __add__(self, other):
        cpy = self.copy()
        cpy += other
        return cpy


    def __radd__(self, other):
        return self + other


    def __sub__(self, other):
        cpy = self.copy()
        cpy -= other
        return cpy


    def __rsub__(self, other):
        return -self + other


    def __mul__(self, other):
        cpy = self.copy()
        cpy *= other
        return cpy


    def __rmul__(self, other):
        return self * other


    def __truediv__(self, other):
        cpy = self.copy()
        cpy /= other
        return cpy


    def inv(self):
        return 1 / self


    def __rtruediv__(self, other):
        return self.inv() * other


    def __pow__(self, other):
        cpy = self.copy()
        cpy **= other
        return cpy


    def __mod__(self, other):
        cpy = self.copy()
        cpy %= other
        return cpy


    def __floordiv__(self, other):
        cpy = self.copy()
        cpy //= other
        return cpy


    def __lshift__(self, other):
        cpy = self.copy()
        cpy <<= other
        return cpy


    def __rshift__(self, other):
        cpy = self.copy()
        cpy >>= other
        return cpy


    def __and__(self, other):
        cpy = self.copy()
        cpy &= other
        return cpy


    def __or__(self, other):
        cpy = self.copy()
        cpy |= other
        return cpy


    def __xor__(self, other):
        cpy = self.copy()
        cpy ^= other
        return cpy


    def op(self, other, operator=operator.add):
        cpy = self.copy()
        cpy.iop(other, operator)   # not implemented in this class
        return cpy




class BaseFunction(Arithmetic):
    # function without domain
    # func: function or number (regarded as constant function) or set (as characteristic function)
    def __init__(self, func):
        if isinstance(func, BaseFunction):   # deprecated
            self.func = func.func
            # print('It is deprecated to set func as a Function object.')
        elif isinstance(func, FunctionType):
            self.func = func
        elif isinstance(func, set):           # deprecated also
            self.func = lambda x: x in func
        else:           # func is a number
            self.func = func


    def copy(self):
        return BaseFunction(self.func)


    def __call__(self, *args, **kwargs):
        # called as a function
        if isinstance(self.func, FunctionType):
            return self.func(*args, **kwargs)
        else:            # func is a number
            return self.func


    # general operator:
    def iop(self, other, operator=operator.add):
        # augmented assignment
        cpy = self.copy()       # copy it neccessarily
        if isinstance(other, BaseFunction):
            def f(*args, **kwargs):
                return operator(cpy(*args, **kwargs), other(*args, **kwargs))
        elif isinstance(other, FunctionType):
            def f(*args, **kwargs):
                return operator(cpy(*args, **kwargs), other(*args, **kwargs))
        else:
            def f(*args, **kwargs):
                return operator(cpy(*args, **kwargs), other)
        self.func = f
        return self


    def rop(self, other, operator=operator.add):
        # right operator
        cpy = self.copy()
        if isinstance(other, BaseFunction):
            def f(*args, **kwargs):
                return operator(other(*args, **kwargs), cpy(*args, **kwargs))
        elif isinstance(other, FunctionType):
            def f(*args, **kwargs):
                return operator(other(*args, **kwargs), cpy(*args, **kwargs))
        else:
            def f(*args, **kwargs):
                return operator(other, cpy(*args, **kwargs))
        return self.__class__(f)


    def op1(self, operator=operator.pos):
        # unary operator
        cpy = self.copy()
        def f(*args, **kwargs):
            return operator(cpy(*args, **kwargs))
        return self.__class__(f)


    # operators
    def __iadd__(self, other):
        self.iop(other)
        return self


    def __pos__(self):
        return f.copy()


    def __neg__(self):
        return self.op1(operator=operator.neg)


    def __isub__(self, other):
        self.iop(other, operator=operator.sub)
        return self


    def __imul__(self, other):
        self.iop(other, operator=operator.mul)
        return self


    def __itruediv__(self, other):
        self.iop(other, operator=operator.truediv)
        return self


    def __rtruediv__(self, other):
        return self.rop(other, operator=operator.truediv)


    def __ifloordiv__(self, other):
        self.iop(other, operator=operator.floordiv)
        return self


    def __rfloordiv__(self, other):
        return self.rop(other, operator=operator.floordiv)


    def __ipow__(self, other):
        self.iop(other, operator=operator.pow)
        return self


    def __rpow__(self, other):
        return self.rop(other, operator=operator.pow)


    def __imod__(self, other):
        self.iop(other, operator=operator.mod)
        return self


    def __rmod__(self, other):
        return self.rop(other, operator=operator.mod)


    def __irshift__(self, other):
        self.iop(other, operator=operator.rshift)
        return self


    def __rrshift__(self, other):
        return self.rop(other, operator=operator.rshift)


    def __ilshift__(self, other):
        self.iop(other, operator=operator.lshift)
        return self


    def __rlshift__(self, other):
        return self.rop(other, operator=operator.lshift)


    def __iand__(self, other):
        self.iop(other, operator=operator.and_)
        return self


    def __rand__(self, other):
        return self.rop(other, operator=operator.and_)


    def __ior__(self, other):
        self.iop(other, operator=operator.or_)
        return self


    def __ror__(self, other):
        return self.rop(other, operator=operator.or_)


    def __ixor__(self, other):
        self.iop(other, operator=operator.xor)
        return self


    def __rxor__(self, other):
        return self.rop(other, operator=operator.xor)


    def __abs__(self):
        return self.op1(operator=operator.abs)


    def __invert__(self):
        return self.op1(operator=operator.invert)


    # comparison:
    def __eq__(self, other):
        return self.op(other, operator=operator.eq)


    def __ne__(self, other):
        return self.op(other, operator=operator.ne)


    def __lt__(self, other):
        return self.op(other, operator=operator.lt)


    def __gt__(self, other):
        return self.op(other, operator=operator.gt)


    def __le__(self, other):
        return self.op(other, operator=operator.le)


    def __ge__(self, other):
        return self.op(other, operator=operator.ge)


    def __getitem__(self, index):
        def f(*args, **kwargs):
            return self(*args, **kwargs)[index]
        return self.__class__(f, self.domain)


    def tensor(self, other):
        def f(x, y):
            return self(x) * other(y)
        return self.__class__(f)




class Type(BaseFunction):
    '''implementaion of definition domains of functions
    func: * -> {True, False}'''
    def __init__(self, func=True):
        if isinstance(func, type):
            self.func = lambda x: isinstance(x, func)
        else:
            super(Type, self).__init__(func)


    def copy(self):
        return Type(self.func)


    def isempty(self):
        return self.func is False


    def __contains__(self, a):
        if isinstance(a, tuple):
            return self(*a)
        return self(a)


    def __imul__(self, other):
        # Cartesian product
        cpy = self.copy()        # neccessarily
        if isinstance(other, Type) and other.isempty() or isinstance(other, int) and other is False:
            self.func = False
        else:
            def f(x, y):
                return cpy(x) and other(y)
            self.func = f
        return self


class Domain(Type):
    pass




TRUE = Type()        # x is always in TRUE
FALSE = Type(False)  # x is never in FALSE


def Interval(lb=None, ub=None):
    '''lb: [None]lower bound
    ub: [None]upper bound
    where None represents infinity'''
    if lb is None:
        return Type(lambda x: x <= ub)
    elif ub is None:
        return Type(lambda x: lb <= x)
    else:
        return Type(lambda x: lb <= x <= ub)




class Function(BaseFunction):
    '''Function with domain
    func: FunctionType, function/number
    domain: Type[TRUE], definition domain of func
    codomain: Type, range of func
    name: name of function'''
    def __init__(self, func, domain=None, codomain=None, name=''):
        if isinstance(func, BaseFunction):   # deprecated
            self.func = func.func
            if isinstance(func, Function):
                self.domain = func.domain
            else:
                self.domain = TRUE
        else:
            super(Function, self).__init__(func)
            if domain is None:
                self.domain = TRUE
        if domain:
            self.domain = domain
        self.name = name
        self.codomain = codomain


    def __call__(self, *args, **kwargs):
        # test whether arguments are in the domain, then called as a function
        if not self.domain(*args):
            raise Exception('argument is not in definition domain of the function!')
        if isinstance(self.func, FunctionType):
            # if not self.codomain(ret): ...
            return self.func(*args, **kwargs)
        else:            # func is a number
            return self.func


    def copy(self):
        return Function(self.func, self.domain.copy(), self.name)


    def __repr__(self):
        if self.domain == TRUE:
            return repr(self.func)
        else:
            return '%s: %s'%(repr(self.func), repr(self.domain))


    def reduce(self, *others, operator=operator.add):
        return functools.reduce(lambda f,g: f.op(g, operator), (self,) + others)


    def sum(self, *others):
        return self.reduce(*others)


    def prod(self, *others):
        return self.reduce(*others, operator=operator.mul)


    def restrict(self, dom):
        # restricted on dom
        cpy = self.copy()
        cpy.domain &= dom
        return cpy


    def __or__(self, other):
        if isinstance(other, Type):
            return self.restrict(other)
        return super(Function, self).__or__(other)


    # advanced methods:
    def glue(self, *others, default=None):
        # glue funcitons
        cpy = self.copy()
        if len(others) == 0:
            def f(*args, **kwargs):
                try:
                    return cpy(*args, **kwargs)
                except Exception as ex:                     # not in domain of f and g
                    if default is not None:
                        return default
                    else:
                        raise ex
            return Function(f, self.domain)
        elif len(others) == 1:
            def h(*args, **kwargs):
                try:
                    return cpy(*args, **kwargs)
                except:
                    try:                                    # not in domain of f
                        return others[0](*args, **kwargs)
                    except Exception as ex:                 # not in domain of f and g
                        if default is not None:
                            return default
                        else:
                            raise ex
            return Function(h, self.domain | Function(others[0]).domain)
        else:
            return self.glue(others[0], default=default).glue(*others[1:], default=default)


    def compose(self, *others):
        # compose: self(other(x))
        cpy = self.copy()
        if len(others) == 0:
            return cpy
        elif len(others) == 1:
            def h(*args, **kwargs):
                return self(others[0](*args, **kwargs))
            return Function(h, self.domain)
        else:
            return self.compose(other[0]).compose(*others[1:])


    def o(self, *others):
        return self.compose(*others)


    def rcompose(self, *others):
        # compose: other(self(x))
        cpy = self.copy()
        if len(others) == 0:
            return cpy
        elif len(others) == 1:
            def h(*args, **kwargs):
                return others[0](self(*args, **kwargs))
            return Function(h, self.domain)
        else:
            return self.rcompose(other[-1]).compose(*others[:-1])


    def partial(self, *args):
        return Function(functools.partial(self.func, *args), self.domain)


    def freduce(self, *others, foperator):
        # call fop
        return Function(functools.reduce(lambda f,g: f.fop(g, foperator), (self,) + others), self.domain)


    def fop(self, other, foperator):
        # foperator should act on functions (instead of the values of functions)
        return Function(foperator(self.func, other.func if isinstance(other, Function) else other), self.domain)


    def map(self, *args):
        return map(self.func, *args)


    # others:
    def cat(self, *others):
        def f(*args, **kwargs):
            return self(*args, **kwargs) + [other(*args, **kwargs) for other in others]
        return Function(f, self.domain)


    def inv(self):
        def f(*args, **kwargs):
            return self.func(*args, **kwargs).inv()
        return Function(f, self.domain)




def glue(*fs, default=None):
    # glue(F1,F2, ...) to glue functions F1 F2, ...
    if len(fs) == 1:
        def f(*args, **kwargs):
            try:
                return fs[0](*args, **kwargs)
            except Exception as ex:                     # not in domain of f and g
                if default is not None:
                    return default
                else:
                    raise Exception('out of domain')
        return Function(f)
    elif len(fs) == 2:
        def f(*args, **kwargs):
            try:
                return fs[0](*args, **kwargs)
            except:
                try:                                    # not in domain of f
                    return fs[1](*args, **kwargs)
                except Exception as ex:                 # not in domain of f and g
                    if default is not None:
                        return default
                    else:
                        raise eException('out of domain')
        return Function(f)
    else:
        return glue(glue(fs[0], fs[1], default=default), fs[2], default=default)




def restrict(domain):
    # this is a decorator, restrict a function on domain
    def F(f):
        return Function(f, domain)
    return F


if __name__ == "__main__":
    f = 3
    g = lambda x,y: 2/x


    t = (Type(lambda x:x<5) & Type(int)) * TRUE  # define type(domain) and functions on it
    G = Function(g, t)
    F = Function(f, t)
    F.sum(F, F)(2,2)


    print(F(1,2))                  # call function


    print(G.glue(F, Function(5))(3,4))
    print(glue(G,F)(3,4))


    ID = Function(lambda x:x)
    print(ID.compose(F)(3,4))


    @restrict(Interval(1,2))        # restricting decorator
    def f(x):
        return x


    print(f(1))
    try:
        print(f(3))
    except Exception as ex:
        print(ex)


    def g(x):
        return x
    G=Function(g)
    print(G(3))
    G= G | Interval(1,2)            # restricting method   
    print(G(3))