<s> """<STR_LIT>""" <EOL> from __future__ import absolute_import , unicode_literals <EOL> import socket <EOL> from contextlib import contextmanager <EOL> from functools import partial <EOL> from itertools import count <EOL> from time import sleep <EOL> from . common import ignore_errors <EOL> from . five import range <EOL> from . messaging import Consumer , Producer <EOL> from . log import get_logger <EOL> from . utils import cached_property , nested <EOL> from . utils . encoding import safe_repr <EOL> from . utils . limits import TokenBucket <EOL> __all__ = [ '<STR_LIT>' ] <EOL> logger = get_logger ( __name__ ) <EOL> debug , info , warn , error = logger . debug , logger . info , logger . warn , logger . error <EOL> W_CONN_LOST = """<STR_LIT>""" <EOL> W_CONN_ERROR = """<STR_LIT>""" <EOL> class ConsumerMixin ( object ) : <EOL> """<STR_LIT>""" <EOL> connect_max_retries = None <EOL> should_stop = False <EOL> def get_consumers ( self , Consumer , channel ) : <EOL> raise NotImplementedError ( '<STR_LIT>' ) <EOL> def on_connection_revived ( self ) : <EOL> pass <EOL> def on_consume_ready ( self , connection , channel , consumers , ** kwargs ) : <EOL> pass <EOL> def on_consume_end ( self , connection , channel ) : <EOL> pass <EOL> def on_iteration ( self ) : <EOL> pass <EOL> def on_decode_error ( self , message , exc ) : <EOL> error ( "<STR_LIT>" , <EOL> exc , message . content_type , message . content_encoding , <EOL> safe_repr ( message . body ) ) <EOL> message . ack ( ) <EOL> def on_connection_error ( self , exc , interval ) : <EOL> warn ( W_CONN_ERROR , interval , exc , exc_info = <NUM_LIT:1> ) <EOL> @ contextmanager <EOL> def extra_context ( self , connection , channel ) : <EOL> yield <EOL> def run ( self , _tokens = <NUM_LIT:1> , ** kwargs ) : <EOL> restart_limit = self . restart_limit <EOL> errors = ( self . connection . connection_errors + <EOL> self . connection . channel_errors ) <EOL> while not self . should_stop : <EOL> try : <EOL> if restart_limit . can_consume ( _tokens ) : <EOL> for _ in self . consume ( limit = None , ** kwargs ) : <EOL> pass <EOL> else : <EOL> sleep ( restart_limit . expected_time ( _tokens ) ) <EOL> except errors : <EOL> warn ( W_CONN_LOST , exc_info = <NUM_LIT:1> ) <EOL> @ contextmanager <EOL> def consumer_context ( self , ** kwargs ) : <EOL> with self . Consumer ( ) as ( connection , channel , consumers ) : <EOL> with self . extra_context ( connection , channel ) : <EOL> self . on_consume_ready ( connection , channel , consumers , ** kwargs ) <EOL> yield connection , channel , consumers <EOL> def consume ( self , limit = None , timeout = None , safety_interval = <NUM_LIT:1> , ** kwargs ) : <EOL> elapsed = <NUM_LIT:0> <EOL> with self . consumer_context ( ** kwargs ) as ( conn , channel , consumers ) : <EOL> for i in limit and range ( limit ) or count ( ) : <EOL> if self . should_stop : <EOL> break <EOL> self . on_iteration ( ) <EOL> try : <EOL> conn . drain_events ( timeout = safety_interval ) <EOL> except socket . timeout : <EOL> conn . heartbeat_check ( ) <EOL> elapsed += safety_interval <EOL> if timeout and elapsed >= timeout : <EOL> raise <EOL> except socket . error : <EOL> if not self . should_stop : <EOL> raise <EOL> else : <EOL> yield <EOL> elapsed = <NUM_LIT:0> <EOL> debug ( '<STR_LIT>' ) <EOL> def maybe_conn_error ( self , fun ) : <EOL> """<STR_LIT>""" <EOL> return ignore_errors ( self , fun ) <EOL> def create_connection ( self ) : <EOL> return self . connection . clone ( ) <EOL> @ contextmanager <EOL> def establish_connection ( self ) : <EOL> with self . create_connection ( ) as conn : <EOL> conn . ensure_connection ( self . on_connection_error , <EOL> self . connect_max_retries ) <EOL> yield conn <EOL> @ contextmanager <EOL> def Consumer ( self ) : <EOL> with self . establish_connection ( ) as conn : <EOL> self . on_connection_revived ( ) <EOL> info ( '<STR_LIT>' , conn . as_uri ( ) ) <EOL> channel = conn . default_channel <EOL> cls = partial ( Consumer , channel , <EOL> on_decode_error = self . on_decode_error ) <EOL> with self . _consume_from ( * self . get_consumers ( cls , channel ) ) as c : <EOL> yield conn , channel , c <EOL> debug ( '<STR_LIT>' ) <EOL> self . on_consume_end ( conn , channel ) <EOL> debug ( '<STR_LIT>' ) <EOL> def _consume_from ( self , * consumers ) : <EOL> return nested ( * consumers ) <EOL> @ cached_property <EOL> def restart_limit ( self ) : <EOL> return TokenBucket ( <NUM_LIT:1> ) <EOL> @ cached_property <EOL> def connection_errors ( self ) : <EOL> return self . connection . connection_errors <EOL> @ cached_property <EOL> def channel_errors ( self ) : <EOL> return self . connection . channel_errors <EOL> class ConsumerProducerMixin ( ConsumerMixin ) : <EOL> """<STR_LIT>""" <EOL> _producer_connection = None <EOL> def on_consume_end ( self , connection , channel ) : <EOL> if self . _producer_connection is not None : <EOL> self . _producer_connection . close ( ) <EOL> self . _producer_connection = None <EOL> @ property <EOL> def producer ( self ) : <EOL> return Producer ( self . producer_connection ) <EOL> @ property <EOL> def producer_connection ( self ) : <EOL> if self . _producer_connection is None : <EOL> conn = self . connection . clone ( ) <EOL> conn . ensure_connection ( self . on_connection_error , <EOL> self . connect_max_retries ) <EOL> self . _producer_connection = conn <EOL> return self . _producer_connection </s>
