Source

wdmmg / wdmmg / model / mongo.py

Full commit
from pymongo import Connection 
from pymongo.database import Database
from bson.dbref import DBRef
from pymongo.objectid import ObjectId

connection = ""
db_name = ""


class classproperty(property):
    def __get__(self, cls, owner):
        return self.fget.__get__(None, owner)()


def dictproperty(name):
    def setter(self, value): self[name] = value
    getter = lambda self: self.get(name)
    return property(getter, setter)


class Base(dict):
    
    context = None
    
    @classproperty
    @classmethod
    def collection_name(cls):
        return cls.__name__.lower()
    
    @classproperty
    @classmethod
    def c(cls):
        return db()[cls.collection_name]

    def __setitem__(self, k, v):
        super(Base, self).__setitem__(k, v)
    
    @classmethod
    def find(cls, *a, **kw):
        kw['as_class'] = cls
        return cls.c.find(*a, **kw)
    
    @classmethod
    def find_one(cls, *a, **kw):
        kw['as_class'] = cls
        return cls.c.find_one(*a, **kw)
    
    @classmethod
    def by_name_or_id(cls, name_or_id):
        return cls.find_one({'$or': [{'name': name_or_id},
                                     {'_id': ObjectId(name_or_id)}]})
    
    def __repr__(self):
        dr = super(Base, self).__repr__()
        return "<%s(%s)>" % (self.collection_name, dr)
    
    def to_ref(self):
        if self.id:
            return DBRef(self.collection_name, self.id)
        
    def to_ref_dict(self):
        d = dict(self.items())
        #if '_id' in d: del d['_id']
        d['ref'] = self.to_ref()
        return d
    
    def explain_key(self, key):
        return Explanation.find_one({'collection': self.collection_name,
                                     'context': self.context,
                                     'key': key})
    
    def explain(self):
        explanations = Explanation.find({'collection': self.collection_name,
                                         'context': self.context,
                                         'key': {'$in': self.keys()}})
        return dict([(e.key, e) for e in explanations])
        
    def to_flat_dict(self, sep='.'):
        """ Flatten down a dictionary with some smartness. """
        def _flatten(orig):
            flat = {}
            for k, v in orig.items():
                if isinstance(v, DBRef):
                    v = v.as_doc().to_dict()
                k = k.replace('$', '')
                if k == '_id': 
                    k = 'id'
                if isinstance(v, dict):
                    if 'name' in v.keys():
                        flat[k] = v['name']
                        del v['name']
                    for sk, sv in _flatten(v).items():
                        flat[k + sep + sk] = sv
                else:
                    flat[k] = v
            return flat
        return _flatten(self)
    

class Explanation(Base):
    
    id = dictproperty('_id')
    coll = dictproperty('collection')
    context = dictproperty('context')
    key = dictproperty('key')
    label = dictproperty('label')
    

def make_connection(host, port):
    global connection
    connection = Connection(host, port)
    connection.document_class = Base

def db():
    return Database(connection, db_name)