glom_expr.py

import glom


class F(object):
    """
    F for filter!

    An F object encapsulates boolean comparisons in a similar
    manner to SQLAlchemy or Pandas columns.

    Given a sequence of objects that you want to filter.

    [Check(F(T.a) == 3 & F(T.b) > 3), default=glom.SKIP]
    """
    def __init__(self, spec, op=None, rhs=None):
        self.spec, self.op, self.rhs = spec, op, rhs

    def __eq__(self, other):
        return F(self, lambda a, b: a == b, other)

    def __and__(self, other):
        return F(self, lambda a, b: a and b, other)

    def __or__(self, other):
        return F(self, lambda a, b: a or b, other)

    def __gt__(self, other):
        return F(self, lambda a, b: a > b, other)

    def __lt__(self, other):
        return F(self, lambda a, b: a < b, other)

    def glomit(self, target, scope):
        # TODO: this evaluation could be made a lot faster
        # by avoiding recursion except at the leaf specs
        lhs = scope[glom.glom](target, self.spec, scope)
        if self.op is None and self.rhs is None:
            return lhs
        rhs = scope[glom.glom](target, self.rhs, scope)
        return self.op(lhs, rhs)


if __name__ == "__main__":
    F('a') > 3
    (F('a') > 3) & (F('b') < 10)
    vals = [
        {'a': i, 'b': (i * 2) % 100, 'c': (i * 3) % 100} for i in range(100)]
    results = glom.glom(
        vals, [glom.Check(
            (F('a') > glom.Literal(3)) & (F('b') < glom.Literal(10)) | (F('c') == glom.Literal(3)),
            default=glom.SKIP)])
    for r in results:
        assert r['a'] > 3 and r['b'] < 10 or r['c'] == 3