Source code for mongokat.document

import base64
import copy
from .utils import dotdict
from uuid import UUID, uuid4
from bson import BSON
from pymongo.errors import OperationFailure
import collections


def _flatten_fetched_fields(fields_arg):
    """ this method takes either a kwargs 'fields', which can be a dict :
    {"_id": False, "store": 1, "url": 1} or a list : ["store", "flag", "url"]
    and returns a tuple : ("store", "flag").
    it HAS to be a tuple, so that it is not updated by the different instances
    """
    if fields_arg is None:
        return None
    if isinstance(fields_arg, dict):
        return tuple(sorted([k for k in list(fields_arg.keys()) if fields_arg[k]]))
    else:
        return tuple(sorted(fields_arg))


[docs]class Document(dict): _initialized_with_doc = False _fetched_fields = None mongokat_collection = None gen_skel = True def __init__(self, doc=None, mongokat_collection=None, fetched_fields=None, gen_skel=None): if mongokat_collection is not None: self.mongokat_collection = mongokat_collection self.collection = self.mongokat_collection.collection if fetched_fields is not None: self._fetched_fields = fetched_fields self._fetched_fields = _flatten_fetched_fields(self._fetched_fields) if gen_skel is not None: self.gen_skel = gen_skel if doc is not None: for k, v in doc.items(): self[k] = v if not self._fetched_fields: self._initialized_with_doc = True if self.gen_skel: self.generate_skeleton() def __str__(self): return "%s(%s)" % (self.__class__.__name__, dict(self)) def __hash__(self): if '_id' in self: value = self['_id'] return value.__hash__() else: raise TypeError("A Document is not hashable if it is not saved. Save the document before hashing it") def __deepcopy__(self, memo={}): obj = self.__class__(doc=copy.deepcopy(dict(self), memo), gen_skel=self.gen_skel, mongokat_collection=self.mongokat_collection, fetched_fields=self._fetched_fields) obj.__dict__ = self.__dict__.copy() return obj @property def b64id(self): """ Returns the document's _id as a base64-encoded string """ return base64.b64encode(self["_id"].binary)
[docs] def ensure_fields(self, fields, force_refetch=False): """ Makes sure we fetched the fields, and populate them if not. """ # We fetched with fields=None, we should have fetched them all if self._fetched_fields is None or self._initialized_with_doc: return if force_refetch: missing_fields = fields else: missing_fields = [f for f in fields if f not in self._fetched_fields] if len(missing_fields) == 0: return if "_id" not in self: raise Exception("Can't ensure_fields because _id is missing") self.refetch_fields(missing_fields)
[docs] def refetch_fields(self, missing_fields): """ Refetches a list of fields from the DB """ db_fields = self.mongokat_collection.find_one({"_id": self["_id"]}, fields={k: 1 for k in missing_fields}) self._fetched_fields += tuple(missing_fields) if not db_fields: return for k, v in db_fields.items(): self[k] = v
[docs] def unset_fields(self, fields): """ Removes this list of fields from both the local object and the DB. """ self.mongokat_collection.update_one({"_id": self["_id"]}, {"$unset": { f: 1 for f in fields }}) for f in fields: if f in self: del self[f]
[docs] def reload(self): """ allow to refresh the document, so after using update(), it could reload its value from the database. Be carreful : reload() will erase all unsaved values. If no _id is set in the document, a KeyError is raised. """ old_doc = self.mongokat_collection.find_one({"_id": self['_id']}, read_use="primary") if not old_doc: raise OperationFailure('Can not reload an unsaved document.' ' %s is not found in the database. Maybe _id was a string and not ObjectId?' % self['_id']) else: for k in list(self.keys()): del self[k] self.update(dotdict(old_doc)) self._initialized_with_doc = False
[docs] def delete(self): """ delete the document from the collection from his _id. """ self.mongokat_collection.remove({'_id': self['_id']})
[docs] def save(self, force=False, uuid=False, **kwargs): """ REPLACES the object in DB. This is forbidden with objects from find() methods unless force=True is given. """ if not self._initialized_with_doc and not force: raise Exception("Cannot save a document not initialized from a Python dict. This might remove fields from the DB!") self._initialized_with_doc = False if '_id' not in self: if uuid: self['_id'] = str("%s-%s" % (self.mongokat_collection.__class__.__name__, uuid4())) return self.mongokat_collection.save(self, **kwargs)
[docs] def save_partial(self, data=None, allow_protected_fields=False, **kwargs): """ Saves just the currently set fields in the database. """ # Backwards compat, deprecated argument if "dotnotation" in kwargs: del kwargs["dotnotation"] if data is None: data = dotdict(self) if "_id" not in data: raise KeyError("_id must be set in order to do a save_partial()") del data["_id"] if len(data) == 0: return if not allow_protected_fields: self.mongokat_collection._check_protected_fields(data) apply_on = dotdict(self) self._initialized_with_doc = False self.mongokat_collection.update_one({"_id": self["_id"]}, {"$set": data}, **kwargs) for k, v in data.items(): apply_on[k] = v self.update(dict(apply_on))
def __generate_skeleton(self, doc, struct, path=""): for key in struct: if type(key) is type: new_key = "$%s" % key.__name__ else: new_key = key new_path = ".".join([path, new_key]).strip('.') # # Automatique generate the skeleton with NoneType # if type(key) is not type and key not in doc: if isinstance(struct[key], dict): if isinstance(struct[key], collections.Callable): doc[key] = struct[key]() else: doc[key] = type(struct[key])() elif struct[key] is dict: doc[key] = {} elif isinstance(struct[key], list): doc[key] = type(struct[key])() # elif isinstance(struct[key], CustomType): # if struct[key].init_type is not None: # doc[key] = struct[key].init_type() # else: # doc[key] = None elif struct[key] is list: doc[key] = [] elif isinstance(struct[key], tuple): doc[key] = [None for _ in range(len(struct[key]))] else: doc[key] = None # # if the value is a dict, we have a another structure to validate # if isinstance(struct[key], dict) and type(key) is not type: self.__generate_skeleton(doc[key], struct[key], new_path)
[docs] def generate_skeleton(self): if self.mongokat_collection.structure is not None: self.__generate_skeleton(self, self.mongokat_collection.structure)
[docs] def get_size(self): """ return the size of the underlying bson object """ try: return len(BSON.encode(self)) except: return None
[docs] def validate(self): """ We do not support validation yet. """ pass