Merge pull request #16 from certego/develop

0.1.5

atlasq/queryset/queryset.py

@@ -14,7 +14,13 @@
 
 class AtlasQuerySet(QuerySet):
     def _clone_into(self, new_qs):
-        copy_props = ("index", "_aggrs_query", "_search_result", "_count")
+        copy_props = (
+            "index",
+            "_aggrs_query",
+            "_search_result",
+            "_count",
+            "_return_objects",
+        )
         qs = super()._clone_into(new_qs)
         for prop in copy_props:
             val = getattr(self, prop)
@@ -29,6 +35,7 @@ def __init__(self, document, collection):
         self._aggrs_query: List[Dict[str, Any]] = None
         self._search_result: CommandCursor = None
         self._count: bool = False
+        self._return_objects: bool = True
 
     def ensure_index(self, user: str, password: str, group_id: str, cluster_name: str):
         db_name = self._document._get_db().name  # pylint: disable=protected-access
@@ -39,6 +46,11 @@ def ensure_index(self, user: str, password: str, group_id: str, cluster_name: st
             user, password, group_id, cluster_name, db_name, collection_name
         )
 
+    def __iter__(self):
+        if not self._return_objects:
+            return iter(self._cursor)
+        return super().__iter__()
+
     @property
     def _aggrs(self):
         # corresponding of _query for us
@@ -55,6 +67,8 @@ def _aggrs(self):
     def _cursor(self):
         if not self._search_result:
             self._search_result = self.aggregate(self._aggrs)
+        if not self._return_objects:
+            self._cursor_obj = self._search_result
         return super()._cursor
 
     @property
@@ -70,6 +84,7 @@ def _query(self):
     def __call__(self, q_obj=None, **query):
         if self.index is None:
             raise AtlasIndexError("Index is not set")
+
         q = AtlasQ(**query)
         if q_obj is not None:
             q &= q_obj
@@ -82,7 +97,6 @@ def _get_projections(self) -> List[Dict[str, Any]]:
             if self._query_obj:
                 return [{"$project": {"meta": "$$SEARCH_META"}}]
             return [{"$count": "count"}]
-
         loaded_fields = self._loaded_fields.as_dict()
         logger.debug(loaded_fields)
         if loaded_fields: