<s> """<STR_LIT>""" <EOL> from __future__ import print_function <EOL> from collections import Iterable , namedtuple <EOL> from riak import RiakError <EOL> from six import string_types , PY2 <EOL> from riak . bucket import RiakBucket <EOL> RiakLink = namedtuple ( "<STR_LIT>" , ( "<STR_LIT>" , "<STR_LIT:key>" , "<STR_LIT>" ) ) <EOL> class RiakMapReduce ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , client ) : <EOL> """<STR_LIT>""" <EOL> self . _client = client <EOL> self . _phases = [ ] <EOL> self . _inputs = [ ] <EOL> self . _key_filters = [ ] <EOL> self . _input_mode = None <EOL> def add ( self , arg1 , arg2 = None , arg3 = None , bucket_type = None ) : <EOL> """<STR_LIT>""" <EOL> from riak . riak_object import RiakObject <EOL> if ( arg2 is None ) and ( arg3 is None ) : <EOL> if isinstance ( arg1 , RiakObject ) : <EOL> return self . add_object ( arg1 ) <EOL> else : <EOL> return self . add_bucket ( arg1 , bucket_type ) <EOL> else : <EOL> return self . add_bucket_key_data ( arg1 , arg2 , arg3 , bucket_type ) <EOL> def add_object ( self , obj ) : <EOL> """<STR_LIT>""" <EOL> return self . add_bucket_key_data ( obj . _bucket . _name , obj . _key , None ) <EOL> def add_bucket_key_data ( self , bucket , key , data , bucket_type = None ) : <EOL> """<STR_LIT>""" <EOL> if self . _input_mode == '<STR_LIT>' : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> elif self . _input_mode == '<STR_LIT>' : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> else : <EOL> if isinstance ( key , Iterable ) and not isinstance ( key , string_types ) : <EOL> if bucket_type is not None : <EOL> for k in key : <EOL> self . _inputs . append ( [ bucket , k , data , bucket_type ] ) <EOL> else : <EOL> for k in key : <EOL> self . _inputs . append ( [ bucket , k , data ] ) <EOL> else : <EOL> if bucket_type is not None : <EOL> self . _inputs . append ( [ bucket , key , data , bucket_type ] ) <EOL> else : <EOL> self . _inputs . append ( [ bucket , key , data ] ) <EOL> return self <EOL> def add_bucket ( self , bucket , bucket_type = None ) : <EOL> """<STR_LIT>""" <EOL> self . _input_mode = '<STR_LIT>' <EOL> if isinstance ( bucket , RiakBucket ) : <EOL> if bucket . bucket_type . is_default ( ) : <EOL> self . _inputs = { '<STR_LIT>' : bucket . name } <EOL> else : <EOL> self . _inputs = { '<STR_LIT>' : [ bucket . bucket_type . name , <EOL> bucket . name ] } <EOL> elif bucket_type is not None and bucket_type != "<STR_LIT:default>" : <EOL> self . _inputs = { '<STR_LIT>' : [ bucket_type , bucket ] } <EOL> else : <EOL> self . _inputs = { '<STR_LIT>' : bucket } <EOL> return self <EOL> def add_key_filters ( self , key_filters ) : <EOL> """<STR_LIT>""" <EOL> if self . _input_mode == '<STR_LIT>' : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> self . _key_filters . extend ( key_filters ) <EOL> return self <EOL> def add_key_filter ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> if self . _input_mode == '<STR_LIT>' : <EOL> raise ValueError ( '<STR_LIT>' ) <EOL> self . _key_filters . append ( args ) <EOL> return self <EOL> def search ( self , index , query ) : <EOL> """<STR_LIT>""" <EOL> self . _input_mode = '<STR_LIT>' <EOL> self . _inputs = { '<STR_LIT>' : index , <EOL> '<STR_LIT:index>' : index , <EOL> '<STR_LIT>' : query } <EOL> return self <EOL> def index ( self , bucket , index , startkey , endkey = None , bucket_type = None ) : <EOL> """<STR_LIT>""" <EOL> self . _input_mode = '<STR_LIT>' <EOL> if endkey is None : <EOL> self . _inputs = { '<STR_LIT>' : bucket , <EOL> '<STR_LIT:index>' : index , <EOL> '<STR_LIT:key>' : startkey } <EOL> else : <EOL> self . _inputs = { '<STR_LIT>' : bucket , <EOL> '<STR_LIT:index>' : index , <EOL> '<STR_LIT:start>' : startkey , <EOL> '<STR_LIT:end>' : endkey } <EOL> if bucket_type is not None : <EOL> self . _inputs [ '<STR_LIT>' ] = [ bucket_type , bucket ] <EOL> return self <EOL> def link ( self , bucket = '<STR_LIT:_>' , tag = '<STR_LIT:_>' , keep = False ) : <EOL> """<STR_LIT>""" <EOL> self . _phases . append ( RiakLinkPhase ( bucket , tag , keep ) ) <EOL> return self <EOL> def map ( self , function , options = None ) : <EOL> """<STR_LIT>""" <EOL> if options is None : <EOL> options = dict ( ) <EOL> if isinstance ( function , list ) : <EOL> language = '<STR_LIT>' <EOL> else : <EOL> language = '<STR_LIT>' <EOL> mr = RiakMapReducePhase ( '<STR_LIT>' , <EOL> function , <EOL> options . get ( '<STR_LIT>' , language ) , <EOL> options . get ( '<STR_LIT>' , False ) , <EOL> options . get ( '<STR_LIT>' , None ) ) <EOL> self . _phases . append ( mr ) <EOL> return self <EOL> def reduce ( self , function , options = None ) : <EOL> """<STR_LIT>""" <EOL> if options is None : <EOL> options = dict ( ) <EOL> if isinstance ( function , list ) : <EOL> language = '<STR_LIT>' <EOL> else : <EOL> language = '<STR_LIT>' <EOL> mr = RiakMapReducePhase ( '<STR_LIT>' , <EOL> function , <EOL> options . get ( '<STR_LIT>' , language ) , <EOL> options . get ( '<STR_LIT>' , False ) , <EOL> options . get ( '<STR_LIT>' , None ) ) <EOL> self . _phases . append ( mr ) <EOL> return self <EOL> def run ( self , timeout = None ) : <EOL> """<STR_LIT>""" <EOL> query , link_results_flag = self . _normalize_query ( ) <EOL> try : <EOL> result = self . _client . mapred ( self . _inputs , query , timeout ) <EOL> except RiakError as e : <EOL> if '<STR_LIT>' in e . value : <EOL> for phase in self . _phases : <EOL> if phase . _language == '<STR_LIT>' : <EOL> if type ( phase . _function ) is str : <EOL> raise RiakError ( '<STR_LIT>' <EOL> '<STR_LIT>' <EOL> '<STR_LIT>' + e . value ) <EOL> raise e <EOL> if not ( link_results_flag or <EOL> isinstance ( self . _phases [ - <NUM_LIT:1> ] , RiakLinkPhase ) ) : <EOL> return result <EOL> if result is None : <EOL> return [ ] <EOL> a = [ ] <EOL> for r in result : <EOL> if ( len ( r ) == <NUM_LIT:2> ) : <EOL> link = ( r [ <NUM_LIT:0> ] , r [ <NUM_LIT:1> ] , None ) <EOL> elif ( len ( r ) == <NUM_LIT:3> ) : <EOL> link = ( r [ <NUM_LIT:0> ] , r [ <NUM_LIT:1> ] , r [ <NUM_LIT:2> ] ) <EOL> a . append ( link ) <EOL> return a <EOL> def stream ( self , timeout = None ) : <EOL> """<STR_LIT>""" <EOL> query , lrf = self . _normalize_query ( ) <EOL> return self . _client . stream_mapred ( self . _inputs , query , timeout ) <EOL> def _normalize_query ( self ) : <EOL> num_phases = len ( self . _phases ) <EOL> if num_phases is <NUM_LIT:0> : <EOL> link_results_flag = True <EOL> else : <EOL> link_results_flag = False <EOL> keep_flag = False <EOL> query = [ ] <EOL> for i in range ( num_phases ) : <EOL> phase = self . _phases [ i ] <EOL> if ( i == ( num_phases - <NUM_LIT:1> ) ) and ( not keep_flag ) : <EOL> phase . _keep = True <EOL> if phase . _keep : <EOL> keep_flag = True <EOL> query . append ( phase . to_array ( ) ) <EOL> if isinstance ( self . _inputs , dict ) and len ( self . _inputs ) == <NUM_LIT:1> : <EOL> if len ( self . _key_filters ) > <NUM_LIT:0> : <EOL> self . _inputs [ '<STR_LIT>' ] = self . _key_filters <EOL> else : <EOL> self . _inputs = self . _inputs [ '<STR_LIT>' ] <EOL> return query , link_results_flag <EOL> def map_values ( self , options = None ) : <EOL> """<STR_LIT>""" <EOL> return self . map ( "<STR_LIT>" , options = options ) <EOL> def map_values_json ( self , options = None ) : <EOL> """<STR_LIT>""" <EOL> return self . map ( "<STR_LIT>" , options = options ) <EOL> def reduce_sum ( self , options = None ) : <EOL> """<STR_LIT>""" <EOL> return self . reduce ( "<STR_LIT>" , options = options ) <EOL> def reduce_min ( self , options = None ) : <EOL> """<STR_LIT>""" <EOL> return self . reduce ( "<STR_LIT>" , options = options ) <EOL> def reduce_max ( self , options = None ) : <EOL> """<STR_LIT>""" <EOL> return self . reduce ( "<STR_LIT>" , options = options ) <EOL> def reduce_sort ( self , js_cmp = None , options = None ) : <EOL> """<STR_LIT>""" <EOL> if options is None : <EOL> options = dict ( ) <EOL> if js_cmp : <EOL> options [ '<STR_LIT>' ] = js_cmp <EOL> return self . reduce ( "<STR_LIT>" , options = options ) <EOL> def reduce_numeric_sort ( self , options = None ) : <EOL> """<STR_LIT>""" <EOL> return self . reduce ( "<STR_LIT>" , options = options ) <EOL> def reduce_limit ( self , limit , options = None ) : <EOL> """<STR_LIT>""" <EOL> if options is None : <EOL> options = dict ( ) <EOL> options [ '<STR_LIT>' ] = limit <EOL> code = """<STR_LIT>""" <EOL> return self . reduce ( code , options = options ) <EOL> def reduce_slice ( self , start , end , options = None ) : <EOL> """<STR_LIT>""" <EOL> if options is None : <EOL> options = dict ( ) <EOL> options [ '<STR_LIT>' ] = [ start , end ] <EOL> return self . reduce ( "<STR_LIT>" , options = options ) <EOL> def filter_not_found ( self , options = None ) : <EOL> """<STR_LIT>""" <EOL> return self . reduce ( "<STR_LIT>" , options = options ) <EOL> class RiakMapReducePhase ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , type , function , language , keep , arg ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> if isinstance ( function , string_types ) and PY2 : <EOL> function = function . encode ( '<STR_LIT:ascii>' ) <EOL> except UnicodeError : <EOL> raise TypeError ( '<STR_LIT>' ) <EOL> self . _type = type <EOL> self . _language = language <EOL> self . _function = function <EOL> self . _keep = keep <EOL> self . _arg = arg <EOL> def to_array ( self ) : <EOL> """<STR_LIT>""" <EOL> stepdef = { '<STR_LIT>' : self . _keep , <EOL> '<STR_LIT>' : self . _language , <EOL> '<STR_LIT>' : self . _arg } <EOL> if self . _language == '<STR_LIT>' : <EOL> if isinstance ( self . _function , list ) : <EOL> stepdef [ '<STR_LIT>' ] = self . _function [ <NUM_LIT:0> ] <EOL> stepdef [ '<STR_LIT:key>' ] = self . _function [ <NUM_LIT:1> ] <EOL> elif isinstance ( self . _function , string_types ) : <EOL> if ( "<STR_LIT:{>" in self . _function ) : <EOL> stepdef [ '<STR_LIT:source>' ] = self . _function <EOL> else : <EOL> stepdef [ '<STR_LIT:name>' ] = self . _function <EOL> elif ( self . _language == '<STR_LIT>' and isinstance ( self . _function , list ) ) : <EOL> stepdef [ '<STR_LIT>' ] = self . _function [ <NUM_LIT:0> ] <EOL> stepdef [ '<STR_LIT>' ] = self . _function [ <NUM_LIT:1> ] <EOL> elif ( self . _language == '<STR_LIT>' and <EOL> isinstance ( self . _function , string_types ) ) : <EOL> stepdef [ '<STR_LIT:source>' ] = self . _function <EOL> return { self . _type : stepdef } <EOL> class RiakLinkPhase ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , bucket , tag , keep ) : <EOL> """<STR_LIT>""" <EOL> self . _bucket = bucket <EOL> self . _tag = tag <EOL> self . _keep = keep <EOL> def to_array ( self ) : <EOL> """<STR_LIT>""" <EOL> stepdef = { '<STR_LIT>' : self . _bucket , <EOL> '<STR_LIT>' : self . _tag , <EOL> '<STR_LIT>' : self . _keep } <EOL> return { '<STR_LIT>' : stepdef } <EOL> class RiakKeyFilter ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> if args : <EOL> self . _filters = [ list ( args ) ] <EOL> else : <EOL> self . _filters = [ ] <EOL> def __add__ ( self , other ) : <EOL> f = RiakKeyFilter ( ) <EOL> f . _filters = self . _filters + other . _filters <EOL> return f <EOL> def _bool_op ( self , op , other ) : <EOL> if ( self . _filters and self . _filters [ <NUM_LIT:0> ] [ <NUM_LIT:0> ] == op ) : <EOL> f = RiakKeyFilter ( ) <EOL> f . _filters . extend ( self . _filters ) <EOL> f . _filters [ <NUM_LIT:0> ] . append ( other . _filters ) <EOL> return f <EOL> return RiakKeyFilter ( op , self . _filters , other . _filters ) <EOL> def __and__ ( self , other ) : <EOL> return self . _bool_op ( "<STR_LIT>" , other ) <EOL> def __or__ ( self , other ) : <EOL> return self . _bool_op ( "<STR_LIT>" , other ) <EOL> def __repr__ ( self ) : <EOL> return str ( self . _filters ) <EOL> def __getattr__ ( self , name ) : <EOL> def function ( * args ) : <EOL> args1 = [ name ] + list ( args ) <EOL> other = RiakKeyFilter ( * args1 ) <EOL> return self + other <EOL> return function <EOL> def __iter__ ( self ) : <EOL> return iter ( self . _filters ) <EOL> class RiakMapReduceChain ( object ) : <EOL> """<STR_LIT>""" <EOL> def add ( self , arg1 , arg2 = None , arg3 = None , bucket_type = None ) : <EOL> """<STR_LIT>""" <EOL> mr = RiakMapReduce ( self ) <EOL> return mr . add ( arg1 , arg2 , arg3 , bucket_type ) <EOL> def search ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> mr = RiakMapReduce ( self ) <EOL> return mr . search ( * args ) <EOL> def index ( self , bucket , index , startkey , endkey = None , bucket_type = None ) : <EOL> """<STR_LIT>""" <EOL> mr = RiakMapReduce ( self ) <EOL> return mr . index ( bucket , index , startkey , endkey , bucket_type ) <EOL> def link ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> mr = RiakMapReduce ( self ) <EOL> return mr . link ( * args ) <EOL> def map ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> mr = RiakMapReduce ( self ) <EOL> return mr . map ( * args ) <EOL> def reduce ( self , * args ) : <EOL> """<STR_LIT>""" <EOL> mr = RiakMapReduce ( self ) <EOL> return mr . reduce ( * args ) </s>
