<s> """<STR_LIT>""" <EOL> import copy <EOL> import six <EOL> from . es import ES <EOL> from . filters import ANDFilter , ORFilter , NotFilter , Filter , TermsFilter , TermFilter , RangeFilter , ExistsFilter <EOL> from . facets import Facet , TermFacet <EOL> from . aggs import Agg , TermsAgg <EOL> from . models import ElasticSearchModel <EOL> from . orm . exceptions import DoesNotExist , MultipleObjectsReturned <EOL> from . query import MatchAllQuery , BoolQuery , FilteredQuery , Search <EOL> from . utils import ESRange <EOL> from . utils . compat import integer_types <EOL> REPR_OUTPUT_SIZE = <NUM_LIT:20> <EOL> def get_es_connection ( es_url , es_kwargs ) : <EOL> if es_url : <EOL> es_kwargs . update ( server = es_url ) <EOL> return ES ( ** es_kwargs ) <EOL> class ESModel ( ElasticSearchModel ) : <EOL> def __init__ ( self , index , type , es_url = None , es_kwargs = { } ) : <EOL> self . _index = index <EOL> self . _type = type <EOL> self . objects = QuerySet ( self , es_url = es_url , es_kwargs = es_kwargs ) <EOL> setattr ( self , "<STR_LIT>" , DoesNotExist ) <EOL> setattr ( self , "<STR_LIT>" , MultipleObjectsReturned ) <EOL> def generate_model ( index , doc_type , es_url = None , es_kwargs = { } ) : <EOL> MyModel = type ( '<STR_LIT>' , ( ElasticSearchModel , ) , { } ) <EOL> setattr ( MyModel , "<STR_LIT>" , QuerySet ( MyModel , index = index , type = doc_type , es_url = es_url , es_kwargs = es_kwargs ) ) <EOL> setattr ( MyModel , "<STR_LIT>" , DoesNotExist ) <EOL> setattr ( MyModel , "<STR_LIT>" , MultipleObjectsReturned ) <EOL> return MyModel <EOL> class QuerySet ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , model = None , using = None , index = None , type = None , es_url = None , es_kwargs = { } ) : <EOL> self . es_url = es_url <EOL> self . es_kwargs = es_kwargs <EOL> if model is None and index and type : <EOL> model = ESModel ( index , type , es_url = self . es_url , es_kwargs = self . es_kwargs ) <EOL> self . model = model <EOL> self . _index = index <EOL> if using : <EOL> self . _index = using <EOL> self . _type = type <EOL> self . _queries = [ ] <EOL> self . _filters = [ ] <EOL> self . _facets = [ ] <EOL> self . _aggs = [ ] <EOL> self . _ordering = [ ] <EOL> self . _fields = [ ] <EOL> self . _size = None <EOL> self . _start = <NUM_LIT:0> <EOL> self . _result_cache = None <EOL> def _clear_ordering ( self ) : <EOL> self . _ordering = [ ] <EOL> @ property <EOL> def index ( self ) : <EOL> if self . _index : <EOL> return self . _index <EOL> return self . model . _index <EOL> @ property <EOL> def type ( self ) : <EOL> if self . _type : <EOL> return self . _type <EOL> return self . model . _type <EOL> def __deepcopy__ ( self , memo ) : <EOL> """<STR_LIT>""" <EOL> obj = self . __class__ ( ) <EOL> for k , v in self . __dict__ . items ( ) : <EOL> if k in ( '<STR_LIT>' , '<STR_LIT>' ) : <EOL> obj . __dict__ [ k ] = None <EOL> else : <EOL> obj . __dict__ [ k ] = copy . deepcopy ( v , memo ) <EOL> return obj <EOL> def __getstate__ ( self ) : <EOL> """<STR_LIT>""" <EOL> len ( self ) <EOL> obj_dict = self . __dict__ . copy ( ) <EOL> obj_dict [ '<STR_LIT>' ] = None <EOL> return obj_dict <EOL> def __repr__ ( self ) : <EOL> data = list ( self [ : REPR_OUTPUT_SIZE + <NUM_LIT:1> ] ) <EOL> if len ( data ) > REPR_OUTPUT_SIZE : <EOL> data [ - <NUM_LIT:1> ] = "<STR_LIT>" <EOL> return repr ( data ) <EOL> def _build_query ( self ) : <EOL> if not self . _queries and not self . _filters : <EOL> return MatchAllQuery ( ) <EOL> query = MatchAllQuery ( ) <EOL> if self . _queries : <EOL> if len ( self . _queries ) == <NUM_LIT:1> : <EOL> query = self . _queries [ <NUM_LIT:0> ] <EOL> else : <EOL> query = BoolQuery ( must = self . _queries ) <EOL> if not self . _filters : <EOL> return query <EOL> if len ( self . _filters ) == <NUM_LIT:1> : <EOL> return FilteredQuery ( query , self . _filters [ <NUM_LIT:0> ] ) <EOL> return FilteredQuery ( query , ANDFilter ( self . _filters ) ) <EOL> def _build_search ( self ) : <EOL> query = Search ( self . _build_query ( ) ) <EOL> if self . _ordering : <EOL> query . sort = self . _ordering <EOL> if self . _facets : <EOL> for facet in self . _facets : <EOL> query . facet . add ( facet ) <EOL> if self . _aggs : <EOL> for agg in self . _aggs : <EOL> query . agg . add ( agg ) <EOL> if self . _start is not None : <EOL> query . start = self . _start <EOL> if self . _size is not None : <EOL> query . size = self . _size <EOL> return query <EOL> def _do_query ( self ) : <EOL> return get_es_connection ( self . es_url , self . es_kwargs ) . search ( self . _build_search ( ) , indices = self . index , doc_types = self . type ) <EOL> def __len__ ( self ) : <EOL> if self . _result_cache is None : <EOL> self . _result_cache = self . _do_query ( ) <EOL> return self . _result_cache . total <EOL> def __iter__ ( self ) : <EOL> if self . _result_cache is None : <EOL> len ( self ) <EOL> return iter ( self . _result_cache ) <EOL> def __bool__ ( self ) : <EOL> return self . __nonzero__ ( ) <EOL> def __nonzero__ ( self ) : <EOL> if self . _result_cache is not None : <EOL> len ( self ) <EOL> return bool ( self . _result_cache . total != <NUM_LIT:0> ) <EOL> try : <EOL> next ( iter ( self ) ) <EOL> except StopIteration : <EOL> return False <EOL> return True <EOL> def __getitem__ ( self , k ) : <EOL> """<STR_LIT>""" <EOL> if not isinstance ( k , ( slice , ) + integer_types ) : <EOL> raise TypeError <EOL> assert ( ( not isinstance ( k , slice ) and ( k >= <NUM_LIT:0> ) ) <EOL> or ( isinstance ( k , slice ) and ( k . start is None or k . start >= <NUM_LIT:0> ) <EOL> and ( k . stop is None or k . stop >= <NUM_LIT:0> ) ) ) , "<STR_LIT>" <EOL> if self . _result_cache is None : <EOL> len ( self ) <EOL> return self . _result_cache . __getitem__ ( k ) <EOL> def __and__ ( self , other ) : <EOL> combined = self . _clone ( ) <EOL> if not other . _filters : <EOL> return combined <EOL> if other . _filters : <EOL> combined . _filters . extend ( other . _filters ) <EOL> return combined <EOL> def __or__ ( self , other ) : <EOL> combined = self . _clone ( ) <EOL> if not other . _filters : <EOL> return other . _clone ( ) <EOL> combined . _filters = ORFilter ( [ combined . _filters , other . _filters ] ) <EOL> return combined <EOL> def size ( self , number ) : <EOL> clone = self . _clone ( ) <EOL> clone . _size = number <EOL> return clone <EOL> def start ( self , number ) : <EOL> clone = self . _clone ( ) <EOL> clone . _start = number <EOL> return clone <EOL> def iterator ( self ) : <EOL> """<STR_LIT>""" <EOL> if not self . _result_cache : <EOL> len ( self ) <EOL> for r in self . _result_cache : <EOL> yield r <EOL> def aggregate ( self , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def count ( self ) : <EOL> """<STR_LIT>""" <EOL> return len ( self ) <EOL> def get ( self , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> clone = self . filter ( * args , ** kwargs ) <EOL> num = len ( clone ) <EOL> if num == <NUM_LIT:1> : <EOL> return clone . _result_cache [ <NUM_LIT:0> ] <EOL> if not num : <EOL> raise self . model . DoesNotExist ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> ( self . model . __class__ . __name__ , kwargs ) ) <EOL> raise self . model . MultipleObjectsReturned ( <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" % <EOL> ( self . model . __class__ . __name__ , num , kwargs ) ) <EOL> def create ( self , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> obj = self . model ( ** kwargs ) <EOL> meta = obj . get_meta ( ) <EOL> meta . connection = get_es_connection ( self . es_url , self . es_kwargs ) <EOL> meta . index = self . index <EOL> meta . type = self . type <EOL> obj . save ( force = True ) <EOL> return obj <EOL> def bulk_create ( self , objs , batch_size = None ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def get_or_create ( self , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> assert kwargs , '<STR_LIT>' <EOL> defaults = kwargs . pop ( '<STR_LIT>' , { } ) <EOL> lookup = kwargs . copy ( ) <EOL> try : <EOL> return self . get ( ** lookup ) , False <EOL> except self . model . DoesNotExist : <EOL> params = dict ( [ ( k , v ) for k , v in kwargs . items ( ) if '<STR_LIT>' not in k ] ) <EOL> params . update ( defaults ) <EOL> obj = self . model ( ** params ) <EOL> meta = obj . get_meta ( ) <EOL> meta . connection = get_es_connection ( self . es_url , self . es_kwargs ) <EOL> meta . index = self . index <EOL> meta . type = self . type <EOL> obj . save ( force = True ) <EOL> return obj , True <EOL> def latest ( self , field_name = None ) : <EOL> """<STR_LIT>""" <EOL> latest_by = field_name or "<STR_LIT>" <EOL> assert bool ( latest_by ) , "<STR_LIT>" <EOL> obj = self . _clone ( ) <EOL> obj . _size = <NUM_LIT:1> <EOL> obj . _clear_ordering ( ) <EOL> obj . _ordering . append ( { latest_by : "<STR_LIT>" } ) <EOL> return obj . get ( ) <EOL> def in_bulk ( self , id_list ) : <EOL> """<STR_LIT>""" <EOL> if not id_list : <EOL> return { } <EOL> qs = self . _clone ( ) <EOL> qs . add_filter ( ( '<STR_LIT>' , id_list ) ) <EOL> qs . _clear_ordering ( force_empty = True ) <EOL> return dict ( [ ( obj . _get_pk_val ( ) , obj ) for obj in qs ] ) <EOL> def delete ( self ) : <EOL> """<STR_LIT>""" <EOL> del_query = self . _clone ( ) <EOL> del_query . _clear_ordering ( ) <EOL> get_es_connection ( self . es_url , self . es_kwargs ) . delete_by_query ( self . index , self . type , self . _build_query ( ) ) <EOL> self . _result_cache = None <EOL> def update ( self , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> query = self . _build_query ( ) <EOL> connection = get_es_connection ( self . es_url , self . es_kwargs ) <EOL> results = connection . search ( query , indices = self . index , doc_types = self . type , <EOL> model = self . model , scan = True ) <EOL> for item in results : <EOL> item . update ( kwargs ) <EOL> item . save ( bulk = True ) <EOL> connection . flush_bulk ( True ) <EOL> self . _result_cache = None <EOL> def exists ( self ) : <EOL> if self . _result_cache is None : <EOL> len ( self ) <EOL> return bool ( self . _result_cache . total != <NUM_LIT:0> ) <EOL> def values ( self , * fields ) : <EOL> search = self . _build_search ( ) <EOL> search . facet . reset ( ) <EOL> search . agg . reset ( ) <EOL> search . fields = fields <EOL> return get_es_connection ( self . es_url , self . es_kwargs ) . search ( search , indices = self . index , doc_types = self . type ) <EOL> def values_list ( self , * fields , ** kwargs ) : <EOL> flat = kwargs . pop ( '<STR_LIT>' , False ) <EOL> if kwargs : <EOL> raise TypeError ( '<STR_LIT>' <EOL> % ( kwargs . keys ( ) , ) ) <EOL> if flat and len ( fields ) > <NUM_LIT:1> : <EOL> raise TypeError ( "<STR_LIT>" ) <EOL> assert fields , "<STR_LIT>" <EOL> search = self . _build_search ( ) <EOL> search . facet . reset ( ) <EOL> search . agg . reset ( ) <EOL> search . fields = fields <EOL> if flat : <EOL> return get_es_connection ( self . es_url , self . es_kwargs ) . search ( search , indices = self . index , doc_types = self . type , <EOL> model = lambda x , y : y . get ( "<STR_LIT>" , { } ) . get ( fields [ <NUM_LIT:0> ] , None ) ) <EOL> return get_es_connection ( self . es_url , self . es_kwargs ) . search ( search , indices = self . index , doc_types = self . type ) <EOL> def dates ( self , field_name , kind , order = '<STR_LIT>' ) : <EOL> """<STR_LIT>""" <EOL> assert kind in ( "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ) , "<STR_LIT>" <EOL> assert order in ( '<STR_LIT>' , '<STR_LIT>' ) , "<STR_LIT>" <EOL> search = self . _build_search ( ) <EOL> search . facet . reset ( ) <EOL> search . facet . add_date_facet ( name = field_name . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) , <EOL> field = field_name , interval = kind ) <EOL> search . agg . reset ( ) <EOL> search . agg . add_date_agg ( name = field_name . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) , <EOL> field = field_name , interval = kind ) <EOL> search . size = <NUM_LIT:0> <EOL> resulset = get_es_connection ( self . es_url , self . es_kwargs ) . search ( search , indices = self . index , doc_types = self . type ) <EOL> resulset . fix_aggs ( ) <EOL> entries = [ ] <EOL> for val in resulset . aggs . get ( field_name . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) ) . get ( "<STR_LIT>" , [ ] ) : <EOL> if "<STR_LIT:time>" in val : <EOL> entries . append ( val [ "<STR_LIT:time>" ] ) <EOL> if order == "<STR_LIT>" : <EOL> return sorted ( entries ) <EOL> return sorted ( entries , reverse = True ) <EOL> def none ( self ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def all ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _clone ( ) <EOL> def filter ( self , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> return self . _filter_or_exclude ( False , * args , ** kwargs ) <EOL> def exclude ( self , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> return self . _filter_or_exclude ( True , * args , ** kwargs ) <EOL> def _filter_or_exclude ( self , negate , * args , ** kwargs ) : <EOL> clone = self . _clone ( ) <EOL> filters = self . _build_filter ( * args , ** kwargs ) <EOL> if negate : <EOL> if len ( filters ) > <NUM_LIT:1> : <EOL> filters = ANDFilter ( filters ) <EOL> else : <EOL> filters = filters [ <NUM_LIT:0> ] <EOL> clone . _filters . append ( NotFilter ( filters ) ) <EOL> else : <EOL> clone . _filters . extend ( filters ) <EOL> return clone <EOL> def _build_inner_filter ( self , field , value ) : <EOL> modifiers = ( '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' , '<STR_LIT>' ) <EOL> if field . endswith ( tuple ( [ '<STR_LIT>' . format ( m ) for m in modifiers ] ) ) : <EOL> field , modifier = field . rsplit ( "<STR_LIT>" , <NUM_LIT:1> ) <EOL> else : <EOL> modifier = "<STR_LIT>" <EOL> field = field . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) <EOL> if not modifier or modifier == '<STR_LIT>' : <EOL> if isinstance ( value , list ) : <EOL> return TermsFilter ( field , value ) <EOL> return TermFilter ( field , value ) <EOL> if modifier == "<STR_LIT>" : <EOL> return TermsFilter ( field , value ) <EOL> elif modifier == "<STR_LIT>" : <EOL> return RangeFilter ( ESRange ( field , from_value = value , include_lower = False ) ) <EOL> elif modifier == "<STR_LIT>" : <EOL> return RangeFilter ( ESRange ( field , from_value = value , include_lower = True ) ) <EOL> elif modifier == "<STR_LIT>" : <EOL> return RangeFilter ( ESRange ( field , to_value = value , include_upper = True ) ) <EOL> elif modifier == "<STR_LIT>" : <EOL> return RangeFilter ( ESRange ( field , to_value = value , include_upper = False ) ) <EOL> elif modifier == "<STR_LIT>" : <EOL> return TermsFilter ( field , values = value ) <EOL> elif modifier == "<STR_LIT>" : <EOL> return NotFilter ( TermFilter ( field , value ) ) <EOL> elif modifier == "<STR_LIT>" : <EOL> if isinstance ( value , bool ) and value : <EOL> return ExistsFilter ( field ) <EOL> return NotFilter ( ExistsFilter ( field ) ) <EOL> raise NotImplementedError ( ) <EOL> def _build_filter ( self , * args , ** kwargs ) : <EOL> filters = [ ] <EOL> if args : <EOL> for f in args : <EOL> if isinstance ( f , Filter ) : <EOL> filters . append ( f ) <EOL> else : <EOL> raise TypeError ( '<STR_LIT>' ) <EOL> for field , value in kwargs . items ( ) : <EOL> filters . append ( self . _build_inner_filter ( field , value ) ) <EOL> return filters <EOL> def complex_filter ( self , filter_obj ) : <EOL> """<STR_LIT>""" <EOL> if isinstance ( filter_obj , Filter ) : <EOL> clone = self . _clone ( ) <EOL> clone . _filters . add ( filter_obj ) <EOL> return clone <EOL> return self . _filter_or_exclude ( None , ** filter_obj ) <EOL> def facet ( self , * args , ** kwargs ) : <EOL> return self . annotate ( * args , ** kwargs ) <EOL> def agg ( self , * args , ** kwargs ) : <EOL> return self . annotate ( * args , ** kwargs ) <EOL> def annotate ( self , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> obj = self . _clone ( ) <EOL> for arg in args : <EOL> if isinstance ( arg , Facet ) : <EOL> obj . _facets . append ( arg ) <EOL> elif isinstance ( arg , six . string_types ) : <EOL> obj . _facets . append ( TermFacet ( arg . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) ) ) <EOL> elif isinstance ( arg , Agg ) : <EOL> obj . _aggs . append ( arg ) <EOL> elif isinstance ( arg , six . string_types ) : <EOL> obj . _facets . append ( TermFacet ( arg . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) ) ) <EOL> else : <EOL> raise NotImplementedError ( "<STR_LIT>" ) <EOL> for name , field in kwargs . items ( ) : <EOL> obj . _facets . append ( TermFacet ( field = field . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) , name = name . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) ) ) <EOL> obj . _aggs . append ( TermsAgg ( field = field . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) , name = name . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) ) ) <EOL> return obj <EOL> def order_by ( self , * field_names ) : <EOL> """<STR_LIT>""" <EOL> obj = self . _clone ( ) <EOL> obj . _clear_ordering ( ) <EOL> for field in field_names : <EOL> if field . startswith ( "<STR_LIT:->" ) : <EOL> obj . _ordering . append ( { field . lstrip ( "<STR_LIT:->" ) . replace ( "<STR_LIT>" , "<STR_LIT:.>" ) : "<STR_LIT>" } ) <EOL> else : <EOL> obj . _ordering . append ( { field : "<STR_LIT>" } ) <EOL> return obj <EOL> def distinct ( self , * field_names ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def reverse ( self ) : <EOL> """<STR_LIT>""" <EOL> clone = self . _clone ( ) <EOL> assert self . _ordering , "<STR_LIT>" <EOL> ordering = [ ] <EOL> for order in self . _ordering : <EOL> for k , v in order . items ( ) : <EOL> if v == "<STR_LIT>" : <EOL> ordering . append ( { k : "<STR_LIT>" } ) <EOL> else : <EOL> ordering . append ( { k : "<STR_LIT>" } ) <EOL> clone . _ordering = ordering <EOL> return clone <EOL> def defer ( self , * fields ) : <EOL> """<STR_LIT>""" <EOL> raise NotImplementedError ( ) <EOL> def only ( self , * fields ) : <EOL> """<STR_LIT>""" <EOL> clone = self . _clone ( ) <EOL> clone . _fields = fields <EOL> return clone <EOL> def using ( self , alias ) : <EOL> """<STR_LIT>""" <EOL> clone = self . _clone ( ) <EOL> clone . _index = alias <EOL> return clone <EOL> def ordered ( self ) : <EOL> """<STR_LIT>""" <EOL> return len ( self . _ordering ) > <NUM_LIT:0> <EOL> ordered = property ( ordered ) <EOL> @ classmethod <EOL> def from_qs ( cls , qs , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> assert issubclass ( cls , QuerySet ) , "<STR_LIT>" % cls <EOL> assert isinstance ( qs , QuerySet ) , "<STR_LIT>" <EOL> return qs . _clone ( klass = cls , ** kwargs ) <EOL> def evaluated ( self ) : <EOL> """<STR_LIT>""" <EOL> return not self . _result_cache is None <EOL> @ property <EOL> def facets ( self ) : <EOL> if len ( self . _facets ) == <NUM_LIT:0> : <EOL> return { } <EOL> if self . _result_cache is None : <EOL> len ( self ) <EOL> return self . _result_cache . facets <EOL> @ property <EOL> def aggs ( self ) : <EOL> if len ( self . _aggs ) == <NUM_LIT:0> : <EOL> return { } <EOL> if self . _result_cache is None : <EOL> len ( self ) <EOL> return self . _result_cache . aggs <EOL> def _clone ( self , klass = None , setup = False , ** kwargs ) : <EOL> if klass is None : <EOL> klass = self . __class__ <EOL> c = klass ( model = self . model , using = self . index , index = self . index , type = self . type , es_url = self . es_url , es_kwargs = self . es_kwargs ) <EOL> c . __dict__ . update ( kwargs ) <EOL> c . _queries = list ( self . _queries ) <EOL> c . _filters = list ( self . _filters ) <EOL> c . _facets = list ( self . _facets ) <EOL> c . _aggs = list ( self . _aggs ) <EOL> c . _fields = list ( self . _fields ) <EOL> c . _ordering = list ( self . _ordering ) <EOL> c . _size = self . _size <EOL> c . _start = self . _start <EOL> return c <EOL> value_annotation = True </s>
