class ValidationError(Exception): pass class BaseField(object): """A base class for fields in a MongoDB document. Instances of this class may be added to subclasses of `Document` to define a document's schema. """ def __init__(self, name=None, required=False, default=None): self.name = name self.required = required self.default = default def __get__(self, instance, owner): """Descriptor for retrieving a value from a field in a document. Do any necessary conversion between Python and MongoDB types. """ if instance is None: # Document class being used rather than a document object return self # Get value from document instance if available, if not use default value = instance._data.get(self.name) if value is not None: value = self._to_python(value) elif self.default is not None: value = self.default if callable(value): value = value() return value def __set__(self, instance, value): """Descriptor for assigning a value to a field in a document. Do any necessary conversion between Python and MongoDB types. """ if value is not None: try: value = self._to_python(value) self._validate(value) value = self._to_mongo(value) except ValueError: raise ValidationError('Invalid value for field of type "' + self.__class__.__name__ + '"') elif self.required: raise ValidationError('Field "%s" is required' % self.name) instance._data[self.name] = value def _to_python(self, value): """Convert a MongoDB-compatible type to a Python type. """ return unicode(value) def _to_mongo(self, value): """Convert a Python type to a MongoDB-compatible type. """ return self._to_python(value) def _validate(self, value): """Perform validation on a value. """ return value class DocumentMetaclass(type): """Metaclass for all documents. """ def __new__(cls, name, bases, attrs): metaclass = attrs.get('__metaclass__') super_new = super(DocumentMetaclass, cls).__new__ if metaclass and issubclass(metaclass, DocumentMetaclass): return super_new(cls, name, bases, attrs) doc_fields = {} # Include all fields present in superclasses for base in bases: if hasattr(base, '_fields'): doc_fields.update(base._fields) # Add the document's fields to the _fields attribute for attr_name, attr_value in attrs.items(): if issubclass(attr_value.__class__, BaseField): #print attr_value.name if not attr_value.name: attr_value.name = attr_name doc_fields[attr_name] = attr_value attrs['_fields'] = doc_fields return super_new(cls, name, bases, attrs) class TopLevelDocumentMetaclass(DocumentMetaclass): """Metaclass for top-level documents (i.e. documents that have their own collection in the database. """ def __new__(cls, name, bases, attrs): super_new = super(TopLevelDocumentMetaclass, cls).__new__ # Classes defined in this package are abstract and should not have # their own metadata with DB collection, etc. # __metaclass__ is only set on the class with the __metaclass__ # attribute (i.e. it is not set on subclasses). This differentiates # 'real' documents from the 'Document' class if attrs.get('__metaclass__') == TopLevelDocumentMetaclass: return super_new(cls, name, bases, attrs) collection = name.lower() # Subclassed documents inherit collection from superclass for base in bases: if hasattr(base, '_meta') and 'collection' in base._meta: collection = base._meta['collection'] # Get primary key field object_id_field = None for attr_name, attr_value in attrs.items(): if issubclass(attr_value.__class__, BaseField): if hasattr(attr_value, 'object_id') and attr_value.object_id: object_id_field = attr_name attr_value.required = True meta = { 'collection': collection, 'object_id_field': object_id_field, } meta.update(attrs.get('meta', {})) attrs['_meta'] = meta return super_new(cls, name, bases, attrs) class BaseDocument(object): def __init__(self, **values): self._data = {} # Assign initial values to instance for attr_name, attr_value in self._fields.items(): if attr_name in values: setattr(self, attr_name, values.pop(attr_name)) else: if attr_value.required: raise ValidationError('Field "%s" is required' % self.name) # Use default value setattr(self, attr_name, getattr(self, attr_name)) def __iter__(self): # Use _data rather than _fields as iterator only looks at names so # values don't need to be converted to Python types return iter(self._data) def __getitem__(self, name): """Dictionary-style field access, return a field's value if present. """ try: return getattr(self, name) except AttributeError: raise KeyError(name) def __setitem__(self, name, value): """Dictionary-style field access, set a field's value. """ # Ensure that the field exists before settings its value if name not in self._fields: raise KeyError(name) return setattr(self, name, value) def _to_mongo(self): """Return data dictionary ready for use with MongoDB. """ return dict((k, v) for k, v in self._data.items() if v is not None)