mongo-query-builder.py

import pprint

def oper(op):
    def _oper(arg):
        return {op: arg}
    return _oper


ne = oper('$ne')
gt = oper('$gt')
gte = oper('$gte')
lt = oper('$lt')
lte = oper('$lte')
in_ = oper('$in')
not_in = oper('$nin')


def multi_op(op):
    def _multi_op(*args):
        return {op: list(args)}
    return _multi_op


and_ = multi_op('$and')
or_ = multi_op('$or')


def mkfield(f):
    if (hasattr(f, 'key')):
        return f.key
    return str(f)


class Query(object):
    def __init__(self, query=None):
        self.query = query or {}

    def filter(self, query):
        rv = dict(self.query)
        rv.update(query)
        return rv


def mkquery(field, value, op=None):
    return Query().filter({str(field): op(value) if op else value})


class Field(object):
    def __init__(self, *components):
        self.components = components

    def __str__(self):
        return '.'.join(mkfield(component) for component in self.components)

    def __eq__(self, other):
        return mkquery(self, other)

    def __ne__(self, other):
        return mkquery(self, other, ne)

    def __gt__(self, other):
        return mkquery(self, other, gt)

    def __ge__(self, other):
        return mkquery(self, other, gte)

    def __lt__(self, other):
        return mkquery(self, other, lt)

    def __le__(self, other):
        return mkquery(self, other, lte)

    def in_(self, others):
        return mkquery(self, others, in_)

    def not_in(self, others):
        return mkquery(self, others, not_in)


if __name__ == '__main__':

    class DocField(object):
        def __init__(self, key):
            self.key = key

    class Model(object):
        pass

    class SubDoc(Model):
        foo = DocField('f')
        bar = DocField('r')
        baz = DocField('z')

    class TestModel(Model):
        sub = DocField('s')

    SubFoo = Field(TestModel.sub, SubDoc.foo)
    SubBar = Field(TestModel.sub, SubDoc.bar)
    SubBaz = Field(TestModel.sub, SubDoc.baz)

    pprint = pprint.pprint

    pprint(str(SubFoo))
    pprint(SubFoo == 'x')
    pprint(and_(
        SubFoo == 'x',
        SubBar == 'y',
        SubBaz == 'z'
    ))

    pprint(or_(
        SubFoo == 'x',
        SubBar == 'y',
        SubBaz == 'z'
    ))

    pprint(SubFoo.in_(['x', 'y', 'z']))
    pprint(SubFoo.not_in(['x', 'y', 'z']))

    pprint(and_(
        or_(
            SubFoo.in_(['x', 'y', 'z']),
            SubBar.not_in(['x', 'y', 'z']),
            SubBaz > 4,
            SubBaz < 10,
            SubBaz >= 12,
            SubBaz <= 38

        )
    ))