json-selector.py

from collections import deque
import copy

class Selector(object):
    def __init__(self, fields):
        self.model = {}
        self._build(fields)

    def _build(self, fields):
        for expr in fields:
            self._build_from_expr(expr)

    def _build_from_expr(self, expr):
        parts = deque(expr.split('.'))
        self._build_from_array(parts)

    def _build_from_array(self, parts):
        t = self.model
        while True:
            if len(parts) == 0:
                break
            e = parts.popleft()
            if e not in t:
                t[e] = None
            
            if t[e] is None and len(parts) > 0:
                t[e] = {}

            t = t[e]

    def find(self, o):
        r = copy.deepcopy(self.model)
        self._copy(r, o)
        return r

    def iter(self, s):
        for o in s:
            r = copy.deepcopy(self.model)
            self._copy(r, o)
            yield r

    def _copy(self, r, o):
        if isinstance(o, list):
            r = [copy.deepcopy(r) for x in xrange(len(o))]
            self._copy_array(r, o)
        elif isinstance(r, dict):
            self._copy_dict(r, o)
        else:
            r = o
        return r

    def _copy_array(self, r, o):
        for i in xrange(len(o)):
            self._copy(r[i], o[i])

    def _copy_dict(self, r, o):
        for k in r.keys():
            if k not in o:
                del r[k]
            elif r[k] is None:
                r[k] = o[k]
            else:
                r[k] = self._copy(r[k], o[k])
    
    def __repr__(self):
        return self.model

    def __str__(self):
        return str(self.model)