仿照selalchemy实现简单的mongo查询

首先这是一个很奇葩的需求，时间紧迫顺手胡写了一个，以后看看有没有好的思路

def and_(item_list):
    return "%s:[%s]" % ("$and", ','.join(loop_func(item_list)))


def or_(item_list):
    return "%s:[%s]" % ("$or", ','.join(loop_func(item_list)))


def lt(a, b):
    return "%s<%s" % (a.name, b)


def le(a, b):
    return "%s<=%s" % (a.name, b)


def ne(a, b):
    return "%s!=%s" % (a.name, b)


def ge(a, b):
    return "%s>=%s" % (a.name, b)


def eq(a, b):
    return "%s=%s" % (a.name, b)


def gt(a, b):
    return "%s>%s" % (a.name, b)


def in_op(a, b):
    return '%s:{$in:%s}' % (a.name, b)


def notin_op(a, b):
    return '%s:{$nin:%s}' % (a.name, b)


def desc_op(a):
    return a


def asc_op(a):
    return a.asc()


class ColumnOperators(object):
    __slots__ = ['name']

    def __init__(self, name):
        self.name = name

    def __and__(self, other):
        return self.operate(and_, other)

    def __or__(self, other):
        return self.operate(or_, other)

    def __lt__(self, other):
        return self.operate(lt, other)

    def __le__(self, other):
        return self.operate(le, other)

    def __eq__(self, other):
        return self.operate(eq, other)

    def __ne__(self, other):
        return self.operate(ne, other)

    def __gt__(self, other):
        return self.operate(gt, other)

    def __ge__(self, other):
        return self.operate(ge, other)

    def operate(self, op, other):
        return op(self, other)

    def in_(self, other):
        return self.operate(in_op, other)

    def notin_(self, other):
        return self.operate(notin_op, other)

    def desc(self):
        return '{%s:1}' % self.name

    def asc(self):
        return '{%s:-1}' % self.name


dic = {'<': '$lt', '<=': '$lte', '=': '$eq', '>=': '$gte', '>': '$gt', '!=': '$ne'}


def split_func(item, op):
    split_item_list = item.strip().split(op)

    if len(split_item_list) != 2:
        raise NotImplementedError("不支持同时两个及以上比较运算")
    for split_item in split_item_list:
        for key in dic:
            if key in split_item:
                raise NotImplementedError("不支持同时两个及以上比较运算")
    return split_item_list, op


def loop_func(item_list):
    and_or_list = []
    for item in item_list:
        if ">=" in item:
            split_item_list, op = split_func(item, ">=")
            and_or_list.append("{%s:{%s:%s}}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip()))
        elif "<=" in item:
            split_item_list, op = split_func(item, "<=")
            and_or_list.append("{%s:{%s:%s}}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip()))
        elif "!=" in item:
            split_item_list, op = split_func(item, ">=")
            and_or_list.append("{%s:{%s:%s}}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip()))
        elif "<" in item:
            split_item_list, op = split_func(item, "<")
            and_or_list.append("{%s:{%s:%s}}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip()))
        elif "=" in item:
            split_item_list, op = split_func(item, "=")
            and_or_list.append("{%s:{%s:%s}}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip()))
        elif ">" in item:
            split_item_list, op = split_func(item, ">")
            and_or_list.append("{%s:{%s:%s}}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip()))
        else:
            and_or_list.append('{%s}' % item)
    return and_or_list


def handle_op(item):
    if ">=" in item:
        split_item_list, op = split_func(item, ">=")
        return "%s:{%s:%s}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip())
    elif "<=" in item:
        split_item_list, op = split_func(item, "<=")
        return "%s:{%s:%s}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip())
    elif "!=" in item:
        split_item_list, op = split_func(item, "!=")
        return "%s:{%s:%s}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip())
    elif "<" in item:
        split_item_list, op = split_func(item, "<")
        return "%s:{%s:%s}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip())
    elif "=" in item:
        split_item_list, op = split_func(item, "=")
        return "%s:{%s:%s}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip())
    elif ">" in item:
        split_item_list, op = split_func(item, ">")
        return "%s:{%s:%s}" % (split_item_list[0].strip(), dic[op], split_item_list[1].strip())
    else:
        return item


class Query(object):
    __slots__ = ['_projection_list', '_query_criteria_list', '_exec_order']

    def __init__(self, *entities):
        self._projection_list = []
        self._query_criteria_list = []
        self._set_entities(*entities)
        self._exec_order = []

    def _set_entities(self, *entities):
        if entities is not ():
            for ent in list(entities):
                self._projection_list.append(ent.name)

    def filter(self, *criterion):
        if criterion is not ():
            for cri in list(criterion):
                self._query_criteria_list.append(handle_op(cri))
        return self

    def order_by(self, standard):
        if 'sort' not in self._exec_order:
            self._exec_order.append({'sort': standard})
        else:
            raise RuntimeError("sort方法只能调用一次")
        return self

    def limit(self, num):
        if 'limit' not in self._exec_order:
            self._exec_order.append({'limit': num})
        else:
            raise RuntimeError("limit方法只能调用一次")
        return self

    def skip(self, num):
        if 'skip' not in self._exec_order:
            self._exec_order.append({'skip': num})
        else:
            raise RuntimeError("skip方法只能调用一次")
        return self


def query(*args):
    return Query(*args)


class Users(object):
    id = ColumnOperators('id')
    name = ColumnOperators('name')
    age = ColumnOperators('age')


def conditions(self):  # 可以将这个方法写到类中，考虑到尽可能少的暴漏接口，就另外写了
    dic = {}
    if self._projection_list:
        dic["columnStr"] = "{%s}" % (",".join(self._projection_list))
    if self._query_criteria_list:
        if len(self._query_criteria_list) == 1:
            dic["condition"] = "{%s}" % (",".join(self._query_criteria_list))
        else:
            dic["condition"] = "{$and:[%s]}" % (",".join(['{%s}' % con for con in self._query_criteria_list]))
    for i in self._exec_order:
        dic.update(i)
    return dic


if __name__ == '__main__':
    Query.conditions = property(conditions)
    query_obj = query(
        Users.id,
        Users.name
    ).filter(
        and_([Users.id > 5, Users.age > 20]),
        or_([Users.name == 'Tom', Users.name == 'Jade'])
    ).order_by(Users.age.desc()).limit(3)
    print(query_obj.conditions)