fromitertoolsimportchain,dropwhilefromoperatorimportmul,attrgetter,__not__fromdjango.db.models.queryimportREPR_OUTPUT_SIZE,EmptyQuerySetdefmul_it(it1,it2):''' Element-wise iterables multiplications. '''assertlen(it1)==len(it2),\
"Can not element-wise multiply iterables of different length."returnmap(mul,it1,it2)defchain_sing(*iterables_or_items):''' As itertools.chain except that if an argument is not iterable then chain it as a singleton. '''foriter_or_iteminiterables_or_items:ifhasattr(iter_or_item,'__iter__'):foriteminiter_or_item:yielditemelse:yielditer_or_itemclassIableSequence(object):''' Wrapper for sequence of iterable and indexable by non-negative integers objects. That is a sequence of objects which implement __iter__, __len__ and __getitem__ for slices, ints and longs. Note: not a Django-specific class. '''def__init__(self,*args,**kwargs):self.iables=args# wrapped sequenceself._len=None# length cacheself._collapsed=[]# collapsed elements cachedef__len__(self):ifnotself._len:self._len=sum(len(iable)foriableinself.iables)returnself._lendef__iter__(self):returnchain(*self.iables)def__nonzero__(self):try:iter(self).next()exceptStopIteration:returnFalsereturnTruedef_collect(self,start=0,stop=None,step=1):ifnotstop:stop=len(self)sub_iables=[]# collect sub setsit=self.iables.__iter__()try:whilestop>start:i=it.next()i_len=len(i)ifi_len>start:# no problem with 'stop' being too bigsub_iables.append(i[start:stop:step])start=max(0,start-i_len)stop-=i_lenexceptStopIteration:passreturnsub_iablesdef__getitem__(self,key):''' Preserves wrapped indexable sequences. Does not support negative indices. '''# params validationifnotisinstance(key,(slice,int,long)):raiseTypeErrorassert((notisinstance(key,slice)and(key>=0))or(isinstance(key,slice)and(key.startisNoneorkey.start>=0)and(key.stopisNoneorkey.stop>=0))), \
"Negative indexing is not supported."# initializationifisinstance(key,slice):start,stop,step=key.indices(len(self))ret_item=Falseelse:# isinstance(key, (int,long))start,stop,step=key,key+1,1ret_item=True# collect sub setsret_iables=self._collect(start,stop,step)# return the simplest possible answerifnotlen(ret_iables):ifret_item:raiseIndexError("'%s' index out of range"%self.__class__.__name__)return()ifret_item:# we have exactly one query set with exactly one itemassertlen(ret_iables)==1andlen(ret_iables[0])==1returnret_iables[0][0]# otherwise we have more then one item in at least one query setiflen(ret_iables)==1:returnret_iables[0]# Note: this can't be self.__class__ instead of IableSequence; exemplary# cause is that indexing over query sets returns lists so we can not# return QuerySetSequence by default. Some type checking enhancement can# be implemented in subclasses.returnIableSequence(*ret_iables)defcollapse(self,stop=None):''' Collapses sequence into a list. Try to do it effectively with caching. '''ifnotstop:stop=len(self)# if we already calculated sufficient collapse then return itiflen(self._collapsed)>=stop:returnself._collapsed[:stop]# otherwise collapse only the missing partitems=self._collapsedsub_iables=self._collect(len(self._collapsed),stop)forsub_iableinsub_iables:items+=sub_iable# cache new collapsed itemsself._collapsed=itemsreturnself._collapseddef__repr__(self):# get +1 element for the truncation msg if applicableitems=self.collapse(stop=REPR_OUTPUT_SIZE+1)iflen(items)>REPR_OUTPUT_SIZE:items[-1]="...(remaining elements truncated)..."returnrepr(items)classQuerySetSequence(IableSequence):''' Wrapper for the query sets sequence without the restriction on the identity of the base models. '''defcount(self):ifnotself._len:self._len=sum(qs.count()forqsinself.iables)returnself._lendef__len__(self):# override: use DB effective count's instead of len()returnself.count()deforder_by(self,*field_names):''' Returns a list of the QuerySetSequence items with the ordering changed. '''# construct a comparator function based on the field names prefixesreverses=[1]*len(field_names)field_names=list(field_names)foriinrange(len(field_names)):field_name=field_names[i]iffield_name[0]=='-':reverses[i]=-1field_names[i]=field_name[1:]# wanna iterable and attrgetter returns single item if 1 arg suppliedfields_getter=lambdai:chain_sing(attrgetter(*field_names)(i))# comparator gets the first non-zero value of the field comparison# results taking into account reverse order for fields prefixed with '-'comparator=lambdai1,i2:\
dropwhile(__not__,mul_it(map(cmp,fields_getter(i1),fields_getter(i2)),reverses)).next()# return new sorted listreturnsorted(self.collapse(),cmp=comparator)deffilter(self,*args,**kwargs):""" Returns a new QuerySetSequence or instance with the args ANDed to the existing set. QuerySetSequence is simplified thus result actually can be one of: QuerySetSequence, QuerySet, EmptyQuerySet. """returnself._filter_or_exclude(False,*args,**kwargs)defexclude(self,*args,**kwargs):""" Returns a new QuerySetSequence instance with NOT (args) ANDed to the existing set. QuerySetSequence is simplified thus result actually can be one of: QuerySetSequence, QuerySet, EmptyQuerySet. """returnself._filter_or_exclude(True,*args,**kwargs)def_simplify(self,qss=None):''' Returns QuerySetSequence, QuerySet or EmptyQuerySet depending on the contents of items, i.e. at least two non empty QuerySets, exactly one non empty QuerySet and all empty QuerySets respectively. Does not modify original QuerySetSequence. '''not_empty_qss=filter(None,qssifqsselseself.iables)ifnotlen(not_empty_qss):returnEmptyQuerySet()iflen(not_empty_qss)==1:returnnot_empty_qss[0]returnQuerySetSequence(*not_empty_qss)def_filter_or_exclude(self,negate,*args,**kwargs):''' Maps _filter_or_exclude over QuerySet items and simplifies the result. '''# each Query set is cloned separatelyreturnself._simplify(*map(lambdaqs:qs._filter_or_exclude(negate,*args,**kwargs),self.iables))defexists(self):forqsinself.iables:ifqs.exists():returnTruereturnFalse