<s> """<STR_LIT>""" <EOL> import re <EOL> from collections import namedtuple , defaultdict <EOL> import datetime <EOL> from corehq . elastic import SIZE_LIMIT <EOL> MISSING_KEY = None <EOL> class AggregationResult ( object ) : <EOL> def __init__ ( self , raw , aggregation ) : <EOL> self . aggregation = aggregation <EOL> self . raw = raw <EOL> self . result = raw . get ( self . aggregation . name , { } ) <EOL> self . _aggregations = self . aggregation . aggregations <EOL> class Aggregation ( object ) : <EOL> name = None <EOL> type = None <EOL> body = None <EOL> result_class = AggregationResult <EOL> aggregations = None <EOL> def __init__ ( self ) : <EOL> raise NotImplementedError ( ) <EOL> def aggregation ( self , aggregation ) : <EOL> if not self . aggregations : <EOL> self . aggregations = [ ] <EOL> self . aggregations . append ( aggregation ) <EOL> return self <EOL> def assemble ( self ) : <EOL> assembled = { self . type : self . body } <EOL> if self . aggregations : <EOL> assembled [ '<STR_LIT>' ] = { } <EOL> for agg in self . aggregations : <EOL> assembled [ '<STR_LIT>' ] [ agg . name ] = agg . assemble ( ) <EOL> return assembled <EOL> def parse_result ( self , result ) : <EOL> return self . result_class ( result , self ) <EOL> class BucketResult ( AggregationResult ) : <EOL> @ property <EOL> def keys ( self ) : <EOL> return [ b [ '<STR_LIT:key>' ] for b in self . normalized_buckets ] <EOL> @ property <EOL> def buckets ( self ) : <EOL> n_buckets = self . normalized_buckets <EOL> buckets = namedtuple ( '<STR_LIT>' , [ b [ '<STR_LIT:key>' ] for b in n_buckets ] ) <EOL> return buckets ( ** { b [ '<STR_LIT:key>' ] : Bucket ( b , self . _aggregations ) for b in n_buckets } ) <EOL> @ property <EOL> def buckets_dict ( self ) : <EOL> return { b [ '<STR_LIT:key>' ] : Bucket ( b , self . _aggregations ) for b in self . normalized_buckets } <EOL> @ property <EOL> def buckets_list ( self ) : <EOL> return { Bucket ( b , self . _aggregations ) for b in self . normalized_buckets } <EOL> @ property <EOL> def raw_buckets ( self ) : <EOL> return self . result [ '<STR_LIT>' ] <EOL> @ property <EOL> def normalized_buckets ( self ) : <EOL> return self . raw_buckets <EOL> def counts_by_bucket ( self ) : <EOL> return { b [ '<STR_LIT:key>' ] : b [ '<STR_LIT>' ] for b in self . normalized_buckets } <EOL> class MissingResult ( AggregationResult ) : <EOL> @ property <EOL> def bucket ( self ) : <EOL> return Bucket ( self . result , self . _aggregations ) <EOL> class TopHitsResult ( AggregationResult ) : <EOL> @ property <EOL> def raw_hits ( self ) : <EOL> return self . result [ '<STR_LIT>' ] [ '<STR_LIT>' ] <EOL> @ property <EOL> def doc_ids ( self ) : <EOL> """<STR_LIT>""" <EOL> return [ r [ '<STR_LIT>' ] for r in self . raw_hits ] <EOL> @ property <EOL> def hits ( self ) : <EOL> """<STR_LIT>""" <EOL> return [ r [ '<STR_LIT>' ] for r in self . raw_hits ] <EOL> @ property <EOL> def total ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . result [ '<STR_LIT>' ] [ '<STR_LIT>' ] <EOL> class StatsResult ( AggregationResult ) : <EOL> @ property <EOL> def count ( self ) : <EOL> return self . result [ '<STR_LIT:count>' ] <EOL> @ property <EOL> def max ( self ) : <EOL> return self . result [ '<STR_LIT>' ] <EOL> @ property <EOL> def min ( self ) : <EOL> return self . result [ '<STR_LIT>' ] <EOL> @ property <EOL> def avg ( self ) : <EOL> return self . result [ '<STR_LIT>' ] <EOL> class ExtendedStatsResult ( StatsResult ) : <EOL> @ property <EOL> def std_dev ( self ) : <EOL> return self . result [ '<STR_LIT>' ] <EOL> class Bucket ( object ) : <EOL> def __init__ ( self , result , aggregations ) : <EOL> self . result = result <EOL> self . aggregations = aggregations <EOL> @ property <EOL> def key ( self ) : <EOL> return self . result . get ( '<STR_LIT:key>' , MISSING_KEY ) <EOL> @ property <EOL> def doc_count ( self ) : <EOL> return self . result [ '<STR_LIT>' ] <EOL> def __getattr__ ( self , attr ) : <EOL> sub_aggregation = filter ( lambda a : a . name == attr , self . aggregations ) [ <NUM_LIT:0> ] <EOL> if sub_aggregation : <EOL> return sub_aggregation . parse_result ( self . result ) <EOL> def __repr__ ( self ) : <EOL> return "<STR_LIT>" . format ( self . key , self . doc_count ) <EOL> class TermsAggregation ( Aggregation ) : <EOL> """<STR_LIT>""" <EOL> type = "<STR_LIT>" <EOL> result_class = BucketResult <EOL> def __init__ ( self , name , field , size = None ) : <EOL> assert re . match ( r'<STR_LIT>' , name ) , "<STR_LIT>" . format ( name ) <EOL> self . name = name <EOL> self . body = { <EOL> "<STR_LIT>" : field , <EOL> "<STR_LIT:size>" : size if size is not None else SIZE_LIMIT , <EOL> } <EOL> class MissingAggregation ( Aggregation ) : <EOL> """<STR_LIT>""" <EOL> type = "<STR_LIT>" <EOL> result_class = MissingResult <EOL> def __init__ ( self , name , field ) : <EOL> assert re . match ( r'<STR_LIT>' , name ) , "<STR_LIT>" . format ( name ) <EOL> self . name = name <EOL> self . body = { "<STR_LIT>" : field } <EOL> class StatsAggregation ( Aggregation ) : <EOL> """<STR_LIT>""" <EOL> type = "<STR_LIT>" <EOL> result_class = StatsResult <EOL> def __init__ ( self , name , field , script = None ) : <EOL> assert re . match ( r'<STR_LIT>' , name ) , "<STR_LIT>" . format ( name ) <EOL> self . name = name <EOL> self . body = { "<STR_LIT>" : field } <EOL> if script : <EOL> self . body . update ( { '<STR_LIT>' : script } ) <EOL> class ExtendedStatsAggregation ( StatsAggregation ) : <EOL> """<STR_LIT>""" <EOL> type = "<STR_LIT>" <EOL> result_class = ExtendedStatsResult <EOL> class TopHitsAggregation ( Aggregation ) : <EOL> """<STR_LIT>""" <EOL> type = "<STR_LIT>" <EOL> result_class = TopHitsResult <EOL> def __init__ ( self , name , field = None , is_ascending = True , size = <NUM_LIT:1> , include = None ) : <EOL> assert re . match ( r'<STR_LIT>' , name ) , "<STR_LIT>" . format ( name ) <EOL> self . name = name <EOL> self . body = { <EOL> '<STR_LIT:size>' : size , <EOL> } <EOL> if field : <EOL> self . body [ "<STR_LIT>" ] = [ { <EOL> field : { <EOL> "<STR_LIT>" : '<STR_LIT>' if is_ascending else '<STR_LIT>' <EOL> } , <EOL> } ] <EOL> if include : <EOL> self . body [ "<STR_LIT>" ] = { "<STR_LIT>" : include } <EOL> class FilterResult ( AggregationResult ) : <EOL> @ property <EOL> def doc_count ( self ) : <EOL> return self . result [ '<STR_LIT>' ] <EOL> class FilterAggregation ( Aggregation ) : <EOL> """<STR_LIT>""" <EOL> type = "<STR_LIT>" <EOL> result_class = FilterResult <EOL> def __init__ ( self , name , filter ) : <EOL> self . name = name <EOL> self . body = filter <EOL> class FiltersAggregation ( Aggregation ) : <EOL> """<STR_LIT>""" <EOL> type = "<STR_LIT>" <EOL> result_class = BucketResult <EOL> def __init__ ( self , name , filters = None ) : <EOL> self . name = name <EOL> self . body = { "<STR_LIT>" : ( filters or { } ) } <EOL> def add_filter ( self , name , filter ) : <EOL> """<STR_LIT>""" <EOL> self . body [ "<STR_LIT>" ] [ name ] = filter <EOL> return self <EOL> class AggregationRange ( namedtuple ( '<STR_LIT>' , '<STR_LIT>' ) ) : <EOL> """<STR_LIT>""" <EOL> def __new__ ( cls , start = None , end = None , key = None ) : <EOL> assert start or end , "<STR_LIT>" <EOL> return super ( AggregationRange , cls ) . __new__ ( cls , start , end , key ) <EOL> def assemble ( self ) : <EOL> range_ = { } <EOL> for key , attr in { '<STR_LIT>' : '<STR_LIT:start>' , '<STR_LIT:to>' : '<STR_LIT:end>' , '<STR_LIT:key>' : '<STR_LIT:key>' } . items ( ) : <EOL> value = getattr ( self , attr ) <EOL> if value : <EOL> if isinstance ( value , datetime . date ) : <EOL> value = value . isoformat ( ) <EOL> elif not isinstance ( value , basestring ) : <EOL> value = unicode ( value ) <EOL> range_ [ key ] = value <EOL> return range_ <EOL> class RangeResult ( BucketResult ) : <EOL> @ property <EOL> def normalized_buckets ( self ) : <EOL> buckets = self . raw_buckets <EOL> if self . aggregation . keyed : <EOL> def _add_key ( key , bucket ) : <EOL> bucket [ '<STR_LIT:key>' ] = key <EOL> return bucket <EOL> return [ _add_key ( k , b ) for k , b in buckets . items ( ) ] <EOL> else : <EOL> def _add_key ( bucket ) : <EOL> key = '<STR_LIT>' . format ( bucket . get ( '<STR_LIT>' , '<STR_LIT:*>' ) , bucket . get ( '<STR_LIT:to>' , '<STR_LIT:*>' ) ) <EOL> bucket [ '<STR_LIT:key>' ] = key <EOL> return bucket <EOL> return [ _add_key ( b ) for b in buckets ] <EOL> class RangeAggregation ( Aggregation ) : <EOL> """<STR_LIT>""" <EOL> type = "<STR_LIT>" <EOL> result_class = RangeResult <EOL> def __init__ ( self , name , field , ranges = None , keyed = True ) : <EOL> self . keyed = keyed <EOL> self . name = name <EOL> self . body = { <EOL> '<STR_LIT>' : field , <EOL> '<STR_LIT>' : keyed , <EOL> '<STR_LIT>' : [ ] <EOL> } <EOL> if ranges : <EOL> for range_ in ranges : <EOL> self . add_range ( range_ ) <EOL> def add_range ( self , range_ ) : <EOL> if isinstance ( range_ , AggregationRange ) : <EOL> range_ = range_ . assemble ( ) <EOL> if range_ . get ( '<STR_LIT:key>' ) : <EOL> self . body [ '<STR_LIT>' ] = True <EOL> self . body [ "<STR_LIT>" ] . append ( range_ ) <EOL> return self <EOL> class HistogramResult ( BucketResult ) : <EOL> def as_facet_result ( self ) : <EOL> return [ <EOL> { '<STR_LIT:time>' : b . key , '<STR_LIT:count>' : b . doc_count } <EOL> for b in self . buckets_list <EOL> ] <EOL> class DateHistogram ( Aggregation ) : <EOL> """<STR_LIT>""" <EOL> type = "<STR_LIT>" <EOL> result_class = HistogramResult <EOL> def __init__ ( self , name , datefield , interval , timezone = None ) : <EOL> self . name = name <EOL> self . body = { <EOL> '<STR_LIT>' : datefield , <EOL> '<STR_LIT>' : interval , <EOL> } <EOL> if timezone : <EOL> self . body [ '<STR_LIT>' ] = timezone <EOL> AggregationTerm = namedtuple ( '<STR_LIT>' , [ '<STR_LIT:name>' , '<STR_LIT>' ] ) <EOL> class NestedAggregationHelper ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , base_query , terms ) : <EOL> self . base_query = base_query <EOL> self . terms = terms <EOL> def get_data ( self ) : <EOL> previous_term = None <EOL> for name , field in reversed ( self . terms ) : <EOL> term = TermsAggregation ( name , field ) <EOL> if previous_term is not None : <EOL> term = term . aggregation ( previous_term ) <EOL> previous_term = term <EOL> query = self . base_query . aggregation ( term ) <EOL> def _add_terms ( aggregation_bucket , term , remaining_terms , current_counts , current_key = None ) : <EOL> for bucket in getattr ( aggregation_bucket , term . name ) . buckets_list : <EOL> key = ( bucket . key , ) if current_key is None else current_key + ( bucket . key , ) <EOL> if remaining_terms : <EOL> _add_terms ( bucket , remaining_terms [ <NUM_LIT:0> ] , remaining_terms [ <NUM_LIT:1> : ] , current_counts , current_key = key ) <EOL> else : <EOL> current_counts [ key ] += bucket . doc_count <EOL> counts = defaultdict ( lambda : <NUM_LIT:0> ) <EOL> _add_terms ( query . size ( <NUM_LIT:0> ) . run ( ) . aggregations , self . terms [ <NUM_LIT:0> ] , self . terms [ <NUM_LIT:1> : ] , current_counts = counts ) <EOL> return self . _format_counts ( counts ) <EOL> def _format_counts ( self , counts ) : <EOL> row_class = namedtuple ( '<STR_LIT>' , [ term . name for term in self . terms ] + [ '<STR_LIT>' ] ) <EOL> for combined_key , count in counts . items ( ) : <EOL> yield row_class ( * ( combined_key + ( count , ) ) ) </s>
