add qcombine class, to combine multiple models in a single queryset

This commit is contained in:
bkfox 2016-06-20 18:10:59 +02:00
parent b195dd74a0
commit bc87e635c2

76
cms/qcombine.py Normal file
View File

@ -0,0 +1,76 @@
import operator
import itertools
import heapq
from django.db.models.query import QuerySet
class QCombine:
"""
This class helps to combine querysets of different models and lists of
object, and to iter over it.
Notes:
- when working on fields, we assume that they exists on all of them;
- for efficiency, there is no possibility to change order per field;
to do so, do it directly on the querysets
"""
order_fields = None
lists = None
def __init__(self, *lists):
"""
lists: list of querysets that are used to initialize the stuff.
"""
self.lists = lists or []
def filter(self, **kwargs):
for qs in self.lists:
if issubclass(type(qs), QuerySet):
qs.filter(**kwargs)
return self
def exclude(self, **kwargs):
for qs in self.lists:
if issubclass(type(qs), QuerySet):
qs.exclude(**kwargs)
return self
def order_by(self, *fields, reverse = False):
"""
Order using these fields. For compatibility, if there is
at least one fields whose name starts with '-', reverse
the order
"""
for i, field in enumerate(fields):
if field[0] == '-':
reverse = True
fields[i] = field[1:]
self.order_reverse = reverse
self.order_fields = fields
for qs in self.lists:
if issubclass(type(qs), QuerySet):
qs.order_by(*fields)
else:
qs.sort(key = operator.attrgetter(fields),
reverse = reverse)
return self
def __len__(self):
return sum([len(qs) for qs in self.lists])
def __iter__(self):
if self.order_fields:
return heapq.merge(
*self.lists,
key = operator.attrgetter(*self.order_fields),
reverse = self.order_reverse
)
return itertools.chain(self.lists)
def __getitem__(self, k):
if type(k) == slice:
return list(itertools.islice(iter(self), k.start, k.stop, k.step))
return list(itertools.islice(iter(self), k))