"""
.. module: mango
"""
import logging
import collections
import datetime
import uuid
import pymongo
logger = logging.getLogger(__name__)
logger.setLevel(logging.DEBUG)
_DOCNAME_TO_TYPE = dict()
class MangoException(Exception):
"""Base exception for all Mango-specific exceptions."""
pass
[docs]def assign(database, collection = None):
"""
A class decorator that assigns a type of document to a database/collection.
If `collection` is ``None``, Mango generates a :class:`Collection` automatically based on the name of the :class:`Document` class.
If `collection` is passed, it must have been generated from the `database`.
Parameters
----------
database : a :class:`Database`
collection : a :class:`Collection`, optional
"""
def assigner(cls, database = database, collection = collection):
if collection is None:
collection = Collection(database, f'coll_{cls.__name__.lower()}')
if collection.database != database:
raise MangoException(f'Cannot assign document type {cls} to {collection}: the collection must belong to {database}, not {collection.database}')
cls.collection = collection
logger.debug(f'Assigned document type {cls} to collection {collection}')
return cls
return assigner
[docs]def register(cls):
"""
A class decorator that "manually" registers a class for conversion when coming back from the database.
Parameters
----------
cls
The class to register.
"""
_DOCNAME_TO_TYPE[cls.__name__] = cls
return cls
class DocumentMetaClass(type):
def __new__(mcs, *args, **kwargs):
clsobj = super().__new__(mcs, *args, **kwargs)
_DOCNAME_TO_TYPE[clsobj.__name__] = clsobj # register class in the global {class name: class} dictionary
return clsobj
[docs]class Document(dict, metaclass = DocumentMetaClass):
"""
The base class for any user-defined documents.
Document is a subclass of :class:`dict`, with dot access to its keys.
There are two forbidden keys: ``'_id'`` and ``'_docname'``.
These are used internally to uniquely identify objects and to reconstruct raw data from the database, respectively.
Mango automatically generates a UUID for the Document to use as its ``_id``, as well as a creation and modification timestamp.
Attributes
----------
id : :class:`str`
A UUID, unique to the document, generated by Mango.
id_filter : :class:`dict`
A dictionary: ``{'_id': self.id}``.
docname : :class:`str`
The name of the Document (i.e., the name of the class)
"""
_internal_keys = ['_id', '_docname']
def __init__(self, **kwargs):
super().__init__(**kwargs)
# this section is unfortunate, but I'm not sure how to optimize it or make it more eloquent
if '_id' not in self:
self['_id'] = uuid.uuid4()
if '_docname' not in self:
self['_docname'] = self.__class__.__name__
now = datetime.datetime.utcnow().replace(microsecond = 0)
if 'timestamp_created' not in self:
self['timestamp_created'] = now
if 'timestamp_modified' not in self:
self['timestamp_modified'] = now
@property
def id(self):
return self._id
@id.setter
def id(self, value):
raise MangoException('Document ids are managed by Mango and should not be modified.')
@property
def docname(self):
return self._docname
@docname.setter
def docname(self, value):
raise MangoException('Document docnames are managed by Mango and should not be modified.')
def __eq__(self, other):
return isinstance(other, self.__class__) and super().__eq__(other)
def __hash__(self):
return self.id
def __getattr__(self, item):
"""Override attribute access to also check the dict, allowing dot access for the dict members."""
try:
return self[item]
except KeyError:
super().__getattribute__(item)
def __setattr__(self, key, value):
"""Override attribute access to actually go the dict."""
self['timestamp_modified'] = datetime.datetime.utcnow().replace(microsecond = 0)
self[key] = value
def __str__(self):
return ''.join((repr(self),
' {',
*(f'\n {k}: {v}' for k, v in sorted(self.items()) if k not in self._internal_keys),
'\n}'))
def __repr__(self):
return f'{self.__class__.__name__} [id = {self.id}]'
@property
def id_filter(self):
return {'_id': self.id}
[docs] def save(self):
"""
Save the Document to its assigned collection.
Returns
-------
result : class:`pymongo.results._WriteResult`
A pymongo database write result object (see `<https://api.mongodb.com/python/current/api/pymongo/results.html>`_)
"""
try:
return self.collection.find_one_and_replace(self.id_filter, self, upsert = True)
except AttributeError: # self.collection doesn't exist, most likely because the document hasn't been assigned to a collection using mango.assign()
raise MangoException(f'Document {repr(self)} cannot be saved because it has not been assigned to a collection.')
[docs] def find_matching(self, *keys, match_document_type = True):
"""
Find all documents in the collection that match this document on the values of the given keys.
Parameters
----------
keys : any number of :class:`str`
Found documents will match this document on these keys.
match_document_type : optional, default `True`
If `True`, found documents must match this document's type as well. Equivalent to adding ``'_docname'`` to `keys`.
Returns
-------
cursor : :class:`Cursor`
A :class:`Cursor` representing the results of the query.
"""
keys = set(keys)
if match_document_type:
keys.add('_docname')
return self.collection.find({k: v for k, v in self.items() if k in keys})
[docs] def find_self(self):
"""Find this Document from the database. Shortcut for ``doc.find_matching('_id')[0]``."""
return self.find_matching('_id')[0]
@classmethod
[docs] def find(cls, **filters):
"""Find all documents in the collection of this type, with additional filters given by the kwargs as if they were arguments to a dictionary used as Collection.find(filters)."""
filters['_docname'] = cls.__name__
return cls.collection.find({**filters})
[docs]def save_many(*documents):
"""
Save many :class:`Documents <Document>` to their associated collections.
This function performs a single bulk write operation for each unique collection among the `documents`.
Parameters
----------
documents : any number of :class:`Documents <Document>`
Documents to be saved.
Returns
-------
results_by_collection : :class:`dict`
A dictionary: ``{collection: write_result}`` (:class:`pymongo.results.BulkWriteResult`).
"""
requests_by_collection = collections.defaultdict(list)
for doc in documents:
requests_by_collection[doc.collection].append(pymongo.ReplaceOne(doc.id_filter, doc, upsert = True))
return {collection: collection.bulk_write(requests) for collection, requests in requests_by_collection.items()}
[docs]class Client(pymongo.MongoClient):
"""
Drop-in replacement for :class:`pymongo.mongo_client.MongoClient`.
When databases and collections are created from a :class:`Client`, they are really instances of Mango's :class:`Database` and :class:`Collection` classes.
"""
def __getitem__(self, name):
"""Get a Mango :class:`Database` instead of a `pymongo` :class:`pymongo.database.Database`."""
return Database(self, name)
[docs]class Database(pymongo.database.Database):
"""
Drop-in replacement for :class:`pymongo.database.Database`.
"""
def __getitem__(self, name):
"""Get a Mango :class:`Collection` instead of a`pymongo` :class:`pymongo.collection.Collection`."""
return Collection(self, name)
def __hash__(self):
return hash(repr(self))
[docs]class Collection(pymongo.collection.Collection):
"""
Drop-in replacement for :class:`pymongo.collection.Collection`.
"""
def find(self, *args, **kwargs):
"""Override to use `Mango`'s :class:`Cursor` instead of `pymongo`'s :class:`pymongo.cursor.Cursor`, passing along the document_type."""
return Cursor(self, *args, **kwargs)
def __hash__(self):
return hash(repr(self))
def _doc_to_obj(doc, depth = 0):
"""Turn a document returned from the database as a dictionary into its corresponding :class:`Document` based on the global ``_DOCNAME_TO_TYPE`` dictionary."""
if isinstance(doc, dict) and '_docname' in doc:
doc = _DOCNAME_TO_TYPE[doc['_docname']](**doc)
elif depth == 0:
doc = Document(**doc)
try:
for k, v in doc.items():
doc[k] = _doc_to_obj(v, depth = depth + 1)
except AttributeError: # doc is not a dictionary
pass
return doc
[docs]class Cursor(pymongo.cursor.Cursor):
"""
Drop-in replacement for :class:`pymongo.cursor.Cursor`.
"""
def _clone_base(self):
"""Override clone method in Cursor to build based on ``self.__class__``, not a hard-coded name."""
return self.__class__(self.__collection)
def __next__(self):
"""Wrap calls to __next__ by the document_type."""
return _doc_to_obj(super().__next__())
def __getitem__(self, item):
"""
Wrap calls to ``__getitem__`` so that they return Mango :class:`Documents <Document>'.
Some care must be taken because the ``super()`` call might return a new :class:`Cursor` over a subset of the originally returned documents instead of a single document (i.e., indexing).
In this case, we should simply return the cursor
"""
from_super = super().__getitem__(item)
if isinstance(from_super, self.__class__):
return from_super
else:
return _doc_to_obj(from_super)