from queryset import QuerySet, QuerySetManager import pymongo 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. """ # Fields may have _types inserted into indexes by default _index_with_types = True def __init__(self, name=None, required=False, default=None, unique=False, unique_with=None, primary_key=False, modified=False): self.name = name if not primary_key else '_id' self.required = required or primary_key self.default = default self.unique = bool(unique or unique_with) self.unique_with = unique_with self.primary_key = primary_key self.modified = modified 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 None: value = self.default # Allow callable default values if callable(value): value = value() return value def __set__(self, instance, value): """Descriptor for assigning a value to a field in a document. """ instance._data[self.name] = value self.modified = True def to_python(self, value): """Convert a MongoDB-compatible type to a Python type. """ return value def to_mongo(self, value): """Convert a Python type to a MongoDB-compatible type. """ return self.to_python(value) def prepare_query_value(self, value): """Prepare a value that is being used in a query for PyMongo. """ return value def validate(self, value): """Perform validation on a value. """ pass class ObjectIdField(BaseField): """An field wrapper around MongoDB's ObjectIds. """ def to_python(self, value): return unicode(value) def to_mongo(self, value): if not isinstance(value, pymongo.objectid.ObjectId): try: return pymongo.objectid.ObjectId(str(value)) except Exception, e: raise ValidationError(e.message) return value def prepare_query_value(self, value): return self.to_mongo(value) def validate(self, value): try: pymongo.objectid.ObjectId(str(value)) except: raise ValidationError('Invalid Object ID') 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 = {} class_name = [name] superclasses = {} simple_class = True for base in bases: # Include all fields present in superclasses if hasattr(base, '_fields'): doc_fields.update(base._fields) class_name.append(base._class_name) # Get superclasses from superclass superclasses[base._class_name] = base superclasses.update(base._superclasses) if hasattr(base, '_meta'): # Ensure that the Document class may be subclassed - # inheritance may be disabled to remove dependency on # additional fields _cls and _types if base._meta.get('allow_inheritance', True) == False: raise ValueError('Document %s may not be subclassed' % base.__name__) else: simple_class = False meta = attrs.get('_meta', attrs.get('meta', {})) if 'allow_inheritance' not in meta: meta['allow_inheritance'] = True # Only simple classes - direct subclasses of Document - may set # allow_inheritance to False if not simple_class and not meta['allow_inheritance']: raise ValueError('Only direct subclasses of Document may set ' '"allow_inheritance" to False') attrs['_meta'] = meta attrs['_class_name'] = '.'.join(reversed(class_name)) attrs['_superclasses'] = superclasses # Add the document's fields to the _fields attribute for attr_name, attr_value in attrs.items(): if hasattr(attr_value, "__class__") and \ issubclass(attr_value.__class__, BaseField): 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() id_field = None base_indexes = [] # Subclassed documents inherit collection from superclass for base in bases: if hasattr(base, '_meta') and 'collection' in base._meta: collection = base._meta['collection'] id_field = id_field or base._meta.get('id_field') base_indexes += base._meta.get('indexes', []) meta = { 'collection': collection, 'max_documents': None, 'max_size': None, 'ordering': [], # default ordering applied at runtime 'indexes': [], # indexes to be ensured at runtime 'id_field': id_field, } # Apply document-defined meta options meta.update(attrs.get('meta', {})) attrs['_meta'] = meta # Set up collection manager, needs the class to have fields so use # DocumentMetaclass before instantiating CollectionManager object new_class = super_new(cls, name, bases, attrs) new_class.objects = QuerySetManager() user_indexes = [QuerySet._build_index_spec(new_class, spec) for spec in meta['indexes']] + base_indexes new_class._meta['indexes'] = user_indexes unique_indexes = [] for field_name, field in new_class._fields.items(): # Generate a list of indexes needed by uniqueness constraints if field.unique: field.required = True unique_fields = [field_name] # Add any unique_with fields to the back of the index spec if field.unique_with: if isinstance(field.unique_with, basestring): field.unique_with = [field.unique_with] # Convert unique_with field names to real field names unique_with = [] for other_name in field.unique_with: parts = other_name.split('.') # Lookup real name parts = QuerySet._lookup_field(new_class, parts) name_parts = [part.name for part in parts] unique_with.append('.'.join(name_parts)) # Unique field should be required parts[-1].required = True unique_fields += unique_with # Add the new index to the list index = [(f, pymongo.ASCENDING) for f in unique_fields] unique_indexes.append(index) # Check for custom primary key if field.primary_key: if not new_class._meta['id_field']: new_class._meta['id_field'] = field_name # Make 'Document.id' an alias to the real primary key field new_class.id = field #new_class._fields['id'] = field else: raise ValueError('Cannot override primary key field') new_class._meta['unique_indexes'] = unique_indexes if not new_class._meta['id_field']: new_class._meta['id_field'] = 'id' new_class.id = new_class._fields['id'] = ObjectIdField(name='_id') return new_class class BaseDocument(object): def __init__(self, **values): self._data = {} modified = 'id' in values.keys() # Assign initial values to instance for attr_name, attr_value in self._fields.items(): attr_value.modified = modified if attr_name in values: setattr(self, attr_name, values.pop(attr_name)) else: # Use default value if present value = getattr(self, attr_name, None) setattr(self, attr_name, value) def validate(self): """Ensure that all fields' values are valid and that required fields are present. """ # Get a list of tuples of field names and their current values fields = [(field, getattr(self, name)) for name, field in self._fields.items()] # Ensure that each field is matched to a valid value for field, value in fields: if value is not None: try: field.validate(value) except (ValueError, AttributeError, AssertionError), e: raise ValidationError('Invalid value for field of type "' + field.__class__.__name__ + '"') elif field.required: raise ValidationError('Field "%s" is required' % field.name) @classmethod def _get_subclasses(cls): """Return a dictionary of all subclasses (found recursively). """ try: subclasses = cls.__subclasses__() except: subclasses = cls.__subclasses__(cls) all_subclasses = {} for subclass in subclasses: all_subclasses[subclass._class_name] = subclass all_subclasses.update(subclass._get_subclasses()) return all_subclasses def __iter__(self): return iter(self._fields) def __getitem__(self, name): """Dictionary-style field access, return a field's value if present. """ try: if name in self._fields: return getattr(self, name) except AttributeError: pass 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 __contains__(self, name): try: val = getattr(self, name) return val is not None except AttributeError: return False def __len__(self): return len(self._data) def __repr__(self): try: u = unicode(self) except (UnicodeEncodeError, UnicodeDecodeError): u = '[Bad Unicode data]' return u'<%s: %s>' % (self.__class__.__name__, u) def __str__(self): if hasattr(self, '__unicode__'): return unicode(self).encode('utf-8') return '%s object' % self.__class__.__name__ def to_mongo(self): """Return data dictionary ready for use with MongoDB. """ data = {} for field_name, field in self._fields.items(): value = getattr(self, field_name, None) if value is not None: data[field.name] = field.to_mongo(value) # Only add _cls and _types if allow_inheritance is not False #if not (hasattr(self, '_meta') and # self._meta.get('allow_inheritance', True) == False): ah = True if hasattr(self, '_meta'): ah = self._meta.get('allow_inheritance', True) if ah: data['_cls'] = self._class_name data['_types'] = self._superclasses.keys() + [self._class_name] return data @classmethod def _from_son(cls, son): """Create an instance of a Document (subclass) from a PyMongo SOM. """ # get the class name from the document, falling back to the given # class if unavailable class_name = son.get(u'_cls', cls._class_name) data = dict((str(key), value) for key, value in son.items()) if '_types' in data: del data['_types'] if '_cls' in data: del data['_cls'] # Return correct subclass for document type if class_name != cls._class_name: subclasses = cls._get_subclasses() if class_name not in subclasses: # Type of document is probably more generic than the class # that has been queried to return this SON return None cls = subclasses[class_name] for field_name, field in cls._fields.items(): if field.name in data: data[field_name] = field.to_python(data[field.name]) return cls(**data)