queryset extension comes here

Author: rochacbruno

flask_mongoengine/__init__.py

@@ -6,13 +6,15 @@
 
 import mongoengine
 
-from mongoengine.queryset import MultipleObjectsReturned, DoesNotExist, QuerySet
+from mongoengine.queryset import (
+    MultipleObjectsReturned, DoesNotExist, QuerySet
+)
 from mongoengine.base import ValidationError
 
 from pymongo import uri_parser
 
-from .sessions import *
-from .pagination import *
+from .sessions import *  # noqa
+from .pagination import *  # noqa
 from .json import overide_json_encoder
 from .wtf import WtfBaseField
 
@@ -87,7 +89,6 @@ def _create_connection(conn_settings):
     return mongoengine.connect(conn.pop('db', 'test'), **conn)
 
 
-
 class MongoEngine(object):
 
     def __init__(self, app=None, config=None):
@@ -106,7 +107,7 @@ def init_app(self, app, config=None):
         # Make documents JSON serializable
         overide_json_encoder(app)
 
-        if not 'mongoengine' in app.extensions:
+        if 'mongoengine' not in app.extensions:
             app.extensions['mongoengine'] = {}
 
         if self in app.extensions['mongoengine']:
@@ -147,6 +148,46 @@ class BaseQuerySet(QuerySet):
     A base queryset with handy extras
     """
 
+    def get_or_create(self, write_options=None, auto_save=True,
+                      *q_objs, **query):
+        """Retrieve unique object or create, if it doesn't exist. Returns a
+        tuple of ``(object, created)``, where ``object`` is the retrieved or
+        created object and ``created`` is a boolean specifying whether a new
+        object was created. Raises
+        :class:`~mongoengine.queryset.MultipleObjectsReturned` or
+        `DocumentName.MultipleObjectsReturned` if multiple results are found.
+        A new document will be created if the document doesn't exists; a
+        dictionary of default values for the new document may be provided as a
+        keyword argument called :attr:`defaults`.
+        .. note:: This requires two separate operations and therefore a
+            race condition exists.  Because there are no transactions in
+            mongoDB other approaches should be investigated, to ensure you
+            don't accidently duplicate data when using this method.  This is
+            now scheduled to be removed before 1.0
+        :param write_options: optional extra keyword arguments used if we
+            have to create a new document.
+            Passes any write_options onto :meth:`~mongoengine.Document.save`
+        :param auto_save: if the object is to be saved automatically if
+            not found.
+
+        add to your documents:
+        meta = {'queryset_class': ExtendedQuerySet}
+        """
+        defaults = query.get('defaults', {})
+        if 'defaults' in query:
+            del query['defaults']
+
+        try:
+            doc = self.get(*q_objs, **query)
+            return doc, False
+        except self._document.DoesNotExist:
+            query.update(defaults)
+            doc = self._document(**query)
+
+            if auto_save:
+                doc.save(write_options=write_options)
+            return doc, True
+
     def get_or_404(self, *args, **kwargs):
         try:
             return self.get(*args, **kwargs)

flask_mongoengine/json.py

@@ -3,9 +3,10 @@
 from mongoengine.base import BaseDocument
 try:
     from mongoengine.base import BaseQuerySet
-except ImportError as ie: # support mongoengine < 0.7
+except ImportError as ie:  # support mongoengine < 0.7
     from mongoengine.queryset import QuerySet as BaseQuerySet
 
+
 def _make_encoder(superclass):
     class MongoEngineJSONEncoder(superclass):
         '''

flask_mongoengine/operation_tracker.py

@@ -35,6 +35,7 @@
 removes = []
 response_sizes = []
 
+
 # Wrap helpers._unpack_response for getting response size
 @functools.wraps(_original_methods['_unpack_response'])
 def _unpack_response(response, *args, **kwargs):
@@ -47,6 +48,7 @@ def _unpack_response(response, *args, **kwargs):
     response_sizes.append(sys.getsizeof(response) / 1024.0)
     return result
 
+
 # Wrap Cursor.insert for getting queries
 @functools.wraps(_original_methods['insert'])
 def _insert(collection_self, doc_or_docs, manipulate=True,
@@ -61,7 +63,7 @@ def _insert(collection_self, doc_or_docs, manipulate=True,
     )
     total_time = (time.time() - start_time) * 1000
 
-    __traceback_hide__ = True
+    __traceback_hide__ = True  # noqa
     stack_trace, internal = _tidy_stacktrace()
     inserts.append({
         'document': doc_or_docs,
@@ -73,6 +75,7 @@ def _insert(collection_self, doc_or_docs, manipulate=True,
     })
     return result
 
+
 # Wrap Cursor.update for getting queries
 @functools.wraps(_original_methods['update'])
 def _update(collection_self, spec, document, upsert=False,
@@ -89,7 +92,7 @@ def _update(collection_self, spec, document, upsert=False,
     )
     total_time = (time.time() - start_time) * 1000
 
-    __traceback_hide__ = True
+    __traceback_hide__ = True  # noqa
     stack_trace, internal = _tidy_stacktrace()
     updates.append({
         'document': document,
@@ -104,6 +107,7 @@ def _update(collection_self, spec, document, upsert=False,
     })
     return result
 
+
 # Wrap Cursor.remove for getting queries
 @functools.wraps(_original_methods['remove'])
 def _remove(collection_self, spec_or_id, safe=None, **kwargs):
@@ -116,7 +120,7 @@ def _remove(collection_self, spec_or_id, safe=None, **kwargs):
     )
     total_time = (time.time() - start_time) * 1000
 
-    __traceback_hide__ = True
+    __traceback_hide__ = True  # noqa
     stack_trace, internal = _tidy_stacktrace()
     removes.append({
         'spec_or_id': spec_or_id,
@@ -128,6 +132,7 @@ def _remove(collection_self, spec_or_id, safe=None, **kwargs):
     })
     return result
 
+
 # Wrap Cursor._refresh for getting queries
 @functools.wraps(_original_methods['refresh'])
 def _cursor_refresh(cursor_self):
@@ -169,11 +174,11 @@ def privar(name):
         if snapshot:
             query_son["$snapshot"] = snapshot
 
-        maxScan = privar("max_scan")
+        maxScan = privar("max_scan")  # noqa
         if maxScan:
             query_son["$maxScan"] = maxScan
 
-    __traceback_hide__ = True
+    __traceback_hide__ = True  # noqa
     stack_trace, internal = _tidy_stacktrace()
     query_data = {
         'time': total_time,
@@ -209,6 +214,7 @@ def privar(name):
 
     return result
 
+
 def install_tracker():
     if pymongo.collection.Collection.insert != _insert:
         pymongo.collection.Collection.insert = _insert
@@ -221,6 +227,7 @@ def install_tracker():
     if pymongo.helpers._unpack_response != _unpack_response:
         pymongo.helpers._unpack_response = _unpack_response
 
+
 def uninstall_tracker():
     if pymongo.collection.Collection.insert == _insert:
         pymongo.collection.Collection.insert = _original_methods['insert']
@@ -231,7 +238,9 @@ def uninstall_tracker():
     if pymongo.cursor.Cursor._refresh == _cursor_refresh:
         pymongo.cursor.Cursor._refresh = _original_methods['cursor_refresh']
     if pymongo.helpers._unpack_response == _unpack_response:
-        pymongo.helpers._unpack_response = _original_methods['_unpack_response']
+        pymongo.helpers._unpack_response = _original_methods[
+            '_unpack_response']
+
 
 def reset():
     global queries, inserts, updates, removes, response_sizes
@@ -241,6 +250,7 @@ def reset():
     removes = []
     response_sizes = []
 
+
 def _get_ordering(son):
     """Helper function to extract formatted ordering from dict.
     """
@@ -250,11 +260,13 @@ def fmt(field, direction):
     if '$orderby' in son:
         return ', '.join(fmt(f, d) for f, d in son['$orderby'].items())
 
+
 def _tidy_stacktrace():
     """
     Tidy the stack_trace
     """
-    socketserver_path = os.path.realpath(os.path.dirname(SocketServer.__file__))
+    socketserver_path = os.path.realpath(
+        os.path.dirname(SocketServer.__file__))
     pymongo_path = os.path.realpath(os.path.dirname(pymongo.__file__))
     paths = ['/site-packages/', '/flaskext/', socketserver_path, pymongo_path]
     internal = False

flask_mongoengine/panels.py

@@ -41,9 +41,12 @@ def nav_title(self):
 
     def nav_subtitle(self):
         attrs = ['queries', 'inserts', 'updates', 'removes']
-        ops = sum(sum((1 for o in getattr(operation_tracker, a)
-                         if not o['internal']))
-                         for a in attrs)
+        ops = sum(
+            sum(
+                (1 for o in getattr(operation_tracker, a)
+                 if not o['internal']))
+            for a in attrs
+        )
         total_time = sum(sum(o['time'] for o in getattr(operation_tracker, a))
                          for a in attrs)
         return '{0} operations in {1:.2f}ms'.format(ops, total_time)
@@ -60,5 +63,6 @@ def content(self):
         context['inserts'] = operation_tracker.inserts
         context['updates'] = operation_tracker.updates
         context['removes'] = operation_tracker.removes
-        context['slow_query_limit'] = current_app.config.get('MONGO_DEBUG_PANEL_SLOW_QUERY_LIMIT', 100)
+        context['slow_query_limit'] = current_app.config.get(
+            'MONGO_DEBUG_PANEL_SLOW_QUERY_LIMIT', 100)
         return self.render('panels/mongo-panel.html', context)

setup.py

@@ -24,7 +24,7 @@
 
 setup(
     name='quokka-flask-mongoengine',
-    version='0.7.2',
+    version='0.7.3',
     url='https://github.com/mongoengine/flask-mongoengine',
     license='BSD',
     author='Ross Lawley',