@@ -13,6 +13,19 @@ def raw_aggregate(self, pipeline, using=None):
13
13
return RawQuerySet (pipeline , model = self .model , using = using )
14
14
15
15
16
+ class RawQuerySet (BaseRawQuerySet ):
17
+ def __init__ (self , pipeline , model = None , using = None ):
18
+ super ().__init__ (pipeline , model = model , using = using )
19
+ self .query = RawQuery (pipeline , using = self .db , model = self .model )
20
+ # Override the superclass's columns property which relies on PEP 249's
21
+ # cursor.description. Instead, RawModelIterable will set the columns
22
+ # based on the keys in the first result.
23
+ self .columns = None
24
+
25
+ def iterator (self ):
26
+ yield from RawModelIterable (self )
27
+
28
+
16
29
class RawQuery (BaseRawQuery ):
17
30
def __init__ (self , pipeline , using , model ):
18
31
self .pipeline = pipeline
@@ -28,24 +41,11 @@ def __str__(self):
28
41
return str (self .pipeline )
29
42
30
43
31
- class RawQuerySet (BaseRawQuerySet ):
32
- def __init__ (self , pipeline , model = None , using = None ):
33
- super ().__init__ (pipeline , model = model , using = using )
34
- self .query = RawQuery (pipeline , using = self .db , model = self .model )
35
- # Override the superclass's columns property which relies on PEP 249's
36
- # cursor.description. Instead, RawModelIterable will set the columns
37
- # based on the keys in the first result.
38
- self .columns = None
39
-
40
- def iterator (self ):
41
- yield from RawModelIterable (self )
42
-
43
-
44
44
class RawModelIterable (BaseRawModelIterable ):
45
45
def __iter__ (self ):
46
46
"""
47
- This is mostly copied from the superclass except for the part that sets
48
- self.queryset.columns from the first document .
47
+ This is copied from the superclass except for the part that sets
48
+ self.queryset.columns from the first result .
49
49
"""
50
50
db = self .queryset .db
51
51
query = self .queryset .query
@@ -55,13 +55,13 @@ def __iter__(self):
55
55
try :
56
56
# Get the columns from the first result.
57
57
try :
58
- first_item = next (query_iterator )
58
+ first_result = next (query_iterator )
59
59
except StopIteration :
60
60
# No results.
61
61
return
62
- self .queryset .columns = list (first_item .keys ())
62
+ self .queryset .columns = list (first_result .keys ())
63
63
# Reset the iterator to include the first item.
64
- query_iterator = self ._make_result (chain ([first_item ], query_iterator ))
64
+ query_iterator = self ._make_result (chain ([first_result ], query_iterator ))
65
65
(
66
66
model_init_names ,
67
67
model_init_pos ,
@@ -88,5 +88,9 @@ def __iter__(self):
88
88
query .cursor .close ()
89
89
90
90
def _make_result (self , query ):
91
+ """
92
+ Convert documents (dictionaries) to tuples as expected by the rest
93
+ of __iter__().
94
+ """
91
95
for result in query :
92
- yield list (result .values ())
96
+ yield tuple (result .values ())
0 commit comments