<s> import json <EOL> import riak . pb . messages <EOL> from riak . util import decode_index_value , bytes_to_str <EOL> from riak . client . index_page import CONTINUATION <EOL> from riak . codecs . ttb import TtbCodec <EOL> from six import PY2 <EOL> class PbufStream ( object ) : <EOL> """<STR_LIT>""" <EOL> _expect = None <EOL> def __init__ ( self , transport , codec ) : <EOL> self . finished = False <EOL> self . transport = transport <EOL> self . codec = codec <EOL> self . resource = None <EOL> def __iter__ ( self ) : <EOL> return self <EOL> def next ( self ) : <EOL> if self . finished : <EOL> raise StopIteration <EOL> try : <EOL> resp_code , data = self . transport . _recv_msg ( ) <EOL> self . codec . maybe_riak_error ( resp_code , data ) <EOL> expect = self . _expect <EOL> self . codec . maybe_incorrect_code ( resp_code , expect ) <EOL> resp = self . codec . parse_msg ( expect , data ) <EOL> except : <EOL> self . finished = True <EOL> raise <EOL> if self . _is_done ( resp ) : <EOL> self . finished = True <EOL> return resp <EOL> def __next__ ( self ) : <EOL> return self . next ( ) <EOL> def _is_done ( self , response ) : <EOL> return response . done <EOL> def attach ( self , resource ) : <EOL> self . resource = resource <EOL> def close ( self ) : <EOL> try : <EOL> while self . next ( ) : <EOL> pass <EOL> except StopIteration : <EOL> pass <EOL> self . resource . release ( ) <EOL> class PbufKeyStream ( PbufStream ) : <EOL> """<STR_LIT>""" <EOL> _expect = riak . pb . messages . MSG_CODE_LIST_KEYS_RESP <EOL> def next ( self ) : <EOL> response = super ( PbufKeyStream , self ) . next ( ) <EOL> if response . done and len ( response . keys ) is <NUM_LIT:0> : <EOL> raise StopIteration <EOL> return response . keys <EOL> def __next__ ( self ) : <EOL> return self . next ( ) <EOL> class PbufMapredStream ( PbufStream ) : <EOL> """<STR_LIT>""" <EOL> _expect = riak . pb . messages . MSG_CODE_MAP_RED_RESP <EOL> def next ( self ) : <EOL> response = super ( PbufMapredStream , self ) . next ( ) <EOL> if response . done and not response . HasField ( '<STR_LIT>' ) : <EOL> raise StopIteration <EOL> return response . phase , json . loads ( bytes_to_str ( response . response ) ) <EOL> def __next__ ( self ) : <EOL> return self . next ( ) <EOL> class PbufBucketStream ( PbufStream ) : <EOL> """<STR_LIT>""" <EOL> _expect = riak . pb . messages . MSG_CODE_LIST_BUCKETS_RESP <EOL> def next ( self ) : <EOL> response = super ( PbufBucketStream , self ) . next ( ) <EOL> if response . done and len ( response . buckets ) is <NUM_LIT:0> : <EOL> raise StopIteration <EOL> return response . buckets <EOL> def __next__ ( self ) : <EOL> return self . next ( ) <EOL> class PbufIndexStream ( PbufStream ) : <EOL> """<STR_LIT>""" <EOL> _expect = riak . pb . messages . MSG_CODE_INDEX_RESP <EOL> def __init__ ( self , transport , codec , index , return_terms = False ) : <EOL> super ( PbufIndexStream , self ) . __init__ ( transport , codec ) <EOL> self . index = index <EOL> self . return_terms = return_terms <EOL> def next ( self ) : <EOL> response = super ( PbufIndexStream , self ) . next ( ) <EOL> if response . done and not ( response . keys or <EOL> response . results or <EOL> response . continuation ) : <EOL> raise StopIteration <EOL> if self . return_terms and response . results : <EOL> return [ ( decode_index_value ( self . index , r . key ) , <EOL> bytes_to_str ( r . value ) ) <EOL> for r in response . results ] <EOL> elif response . keys : <EOL> if PY2 : <EOL> return response . keys [ : ] <EOL> else : <EOL> return [ bytes_to_str ( key ) for key in response . keys ] <EOL> elif response . continuation : <EOL> return CONTINUATION ( bytes_to_str ( response . continuation ) ) <EOL> def __next__ ( self ) : <EOL> return self . next ( ) <EOL> class PbufTsKeyStream ( PbufStream , TtbCodec ) : <EOL> """<STR_LIT>""" <EOL> _expect = riak . pb . messages . MSG_CODE_TS_LIST_KEYS_RESP <EOL> def next ( self ) : <EOL> response = super ( PbufTsKeyStream , self ) . next ( ) <EOL> if response . done and len ( response . keys ) is <NUM_LIT:0> : <EOL> raise StopIteration <EOL> keys = [ ] <EOL> for tsrow in response . keys : <EOL> keys . append ( self . codec . decode_timeseries_row ( tsrow ) ) <EOL> return keys <EOL> def __next__ ( self ) : <EOL> return self . next ( ) </s>
