<s> from __future__ import unicode_literals <EOL> import django <EOL> from django . db import models <EOL> from django . db . models . fields . related import OneToOneField , OneToOneRel <EOL> from django . db . models . query import QuerySet <EOL> from django . core . exceptions import ObjectDoesNotExist <EOL> try : <EOL> from django . db . models . constants import LOOKUP_SEP <EOL> from django . utils . six import string_types <EOL> except ImportError : <EOL> from django . db . models . sql . constants import LOOKUP_SEP <EOL> string_types = ( basestring , ) <EOL> class InheritanceQuerySetMixin ( object ) : <EOL> def select_subclasses ( self , * subclasses ) : <EOL> levels = self . _get_maximum_depth ( ) <EOL> calculated_subclasses = self . _get_subclasses_recurse ( <EOL> self . model , levels = levels ) <EOL> if not subclasses : <EOL> subclasses = calculated_subclasses <EOL> else : <EOL> verified_subclasses = [ ] <EOL> for subclass in subclasses : <EOL> if subclass is self . model : <EOL> continue <EOL> if not isinstance ( subclass , string_types ) : <EOL> subclass = self . _get_ancestors_path ( <EOL> subclass , levels = levels ) <EOL> if subclass in calculated_subclasses : <EOL> verified_subclasses . append ( subclass ) <EOL> else : <EOL> raise ValueError ( <EOL> '<STR_LIT>' % ( <EOL> subclass , '<STR_LIT:U+002CU+0020>' . join ( calculated_subclasses ) ) <EOL> ) <EOL> subclasses = verified_subclasses <EOL> previous_select_related = self . query . select_related <EOL> new_qs = self . select_related ( * subclasses ) <EOL> previous_is_dict = isinstance ( previous_select_related , dict ) <EOL> new_is_dict = isinstance ( new_qs . query . select_related , dict ) <EOL> if previous_is_dict and new_is_dict : <EOL> new_qs . query . select_related . update ( previous_select_related ) <EOL> new_qs . subclasses = subclasses <EOL> return new_qs <EOL> def _clone ( self , klass = None , setup = False , ** kwargs ) : <EOL> for name in [ '<STR_LIT>' , '<STR_LIT>' ] : <EOL> if hasattr ( self , name ) : <EOL> kwargs [ name ] = getattr ( self , name ) <EOL> if django . VERSION < ( <NUM_LIT:1> , <NUM_LIT:9> ) : <EOL> kwargs [ '<STR_LIT>' ] = klass <EOL> kwargs [ '<STR_LIT>' ] = setup <EOL> return super ( InheritanceQuerySetMixin , self ) . _clone ( ** kwargs ) <EOL> def annotate ( self , * args , ** kwargs ) : <EOL> qset = super ( InheritanceQuerySetMixin , self ) . annotate ( * args , ** kwargs ) <EOL> qset . _annotated = [ a . default_alias for a in args ] + list ( kwargs . keys ( ) ) <EOL> return qset <EOL> def iterator ( self ) : <EOL> iter = super ( InheritanceQuerySetMixin , self ) . iterator ( ) <EOL> if getattr ( self , '<STR_LIT>' , False ) : <EOL> extras = tuple ( self . query . extra . keys ( ) ) <EOL> subclasses = sorted ( self . subclasses , key = len , reverse = True ) <EOL> for obj in iter : <EOL> sub_obj = None <EOL> for s in subclasses : <EOL> sub_obj = self . _get_sub_obj_recurse ( obj , s ) <EOL> if sub_obj : <EOL> break <EOL> if not sub_obj : <EOL> sub_obj = obj <EOL> if getattr ( self , '<STR_LIT>' , False ) : <EOL> for k in self . _annotated : <EOL> setattr ( sub_obj , k , getattr ( obj , k ) ) <EOL> for k in extras : <EOL> setattr ( sub_obj , k , getattr ( obj , k ) ) <EOL> yield sub_obj <EOL> else : <EOL> for obj in iter : <EOL> yield obj <EOL> def _get_subclasses_recurse ( self , model , levels = None ) : <EOL> """<STR_LIT>""" <EOL> if django . VERSION < ( <NUM_LIT:1> , <NUM_LIT:8> ) : <EOL> related_objects = model . _meta . get_all_related_objects ( ) <EOL> else : <EOL> related_objects = [ <EOL> f for f in model . _meta . get_fields ( ) <EOL> if isinstance ( f , OneToOneRel ) ] <EOL> rels = [ <EOL> rel for rel in related_objects <EOL> if isinstance ( rel . field , OneToOneField ) <EOL> and issubclass ( rel . field . model , model ) <EOL> and model is not rel . field . model <EOL> ] <EOL> subclasses = [ ] <EOL> if levels : <EOL> levels -= <NUM_LIT:1> <EOL> for rel in rels : <EOL> if levels or levels is None : <EOL> for subclass in self . _get_subclasses_recurse ( <EOL> rel . field . model , levels = levels ) : <EOL> subclasses . append ( <EOL> rel . get_accessor_name ( ) + LOOKUP_SEP + subclass ) <EOL> subclasses . append ( rel . get_accessor_name ( ) ) <EOL> return subclasses <EOL> def _get_ancestors_path ( self , model , levels = None ) : <EOL> """<STR_LIT>""" <EOL> if not issubclass ( model , self . model ) : <EOL> raise ValueError ( <EOL> "<STR_LIT>" % ( model , self . model ) ) <EOL> ancestry = [ ] <EOL> parent_link = model . _meta . get_ancestor_link ( self . model ) <EOL> if levels : <EOL> levels -= <NUM_LIT:1> <EOL> while parent_link is not None : <EOL> if django . VERSION < ( <NUM_LIT:1> , <NUM_LIT:8> ) : <EOL> related = parent_link . related <EOL> else : <EOL> related = parent_link . rel <EOL> ancestry . insert ( <NUM_LIT:0> , related . get_accessor_name ( ) ) <EOL> if levels or levels is None : <EOL> if django . VERSION < ( <NUM_LIT:1> , <NUM_LIT:8> ) : <EOL> parent_model = related . parent_model <EOL> else : <EOL> parent_model = related . model <EOL> parent_link = parent_model . _meta . get_ancestor_link ( <EOL> self . model ) <EOL> else : <EOL> parent_link = None <EOL> return LOOKUP_SEP . join ( ancestry ) <EOL> def _get_sub_obj_recurse ( self , obj , s ) : <EOL> rel , _ , s = s . partition ( LOOKUP_SEP ) <EOL> try : <EOL> node = getattr ( obj , rel ) <EOL> except ObjectDoesNotExist : <EOL> return None <EOL> if s : <EOL> child = self . _get_sub_obj_recurse ( node , s ) <EOL> return child <EOL> else : <EOL> return node <EOL> def get_subclass ( self , * args , ** kwargs ) : <EOL> return self . select_subclasses ( ) . get ( * args , ** kwargs ) <EOL> def _get_maximum_depth ( self ) : <EOL> """<STR_LIT>""" <EOL> levels = None <EOL> if django . VERSION < ( <NUM_LIT:1> , <NUM_LIT:6> , <NUM_LIT:0> ) : <EOL> levels = <NUM_LIT:1> <EOL> return levels <EOL> class InheritanceManagerMixin ( object ) : <EOL> use_for_related_fields = True <EOL> def get_queryset ( self ) : <EOL> return InheritanceQuerySet ( self . model ) <EOL> get_query_set = get_queryset <EOL> def select_subclasses ( self , * subclasses ) : <EOL> return self . get_queryset ( ) . select_subclasses ( * subclasses ) <EOL> def get_subclass ( self , * args , ** kwargs ) : <EOL> return self . get_queryset ( ) . get_subclass ( * args , ** kwargs ) <EOL> class InheritanceQuerySet ( InheritanceQuerySetMixin , QuerySet ) : <EOL> pass <EOL> class InheritanceManager ( InheritanceManagerMixin , models . Manager ) : <EOL> pass <EOL> class QueryManagerMixin ( object ) : <EOL> use_for_related_fields = True <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> if args : <EOL> self . _q = args [ <NUM_LIT:0> ] <EOL> else : <EOL> self . _q = models . Q ( ** kwargs ) <EOL> self . _order_by = None <EOL> super ( QueryManagerMixin , self ) . __init__ ( ) <EOL> def order_by ( self , * args ) : <EOL> self . _order_by = args <EOL> return self <EOL> def get_queryset ( self ) : <EOL> try : <EOL> qs = super ( QueryManagerMixin , self ) . get_queryset ( ) . filter ( self . _q ) <EOL> except AttributeError : <EOL> qs = super ( QueryManagerMixin , self ) . get_query_set ( ) . filter ( self . _q ) <EOL> if self . _order_by is not None : <EOL> return qs . order_by ( * self . _order_by ) <EOL> return qs <EOL> get_query_set = get_queryset <EOL> class QueryManager ( QueryManagerMixin , models . Manager ) : <EOL> pass </s>
