<s> from collections import deque <EOL> import logging <EOL> import time <EOL> from twisted . internet import protocol <EOL> from twisted . internet import defer , reactor <EOL> from twisted . python import failure <EOL> from twimp import amf0 <EOL> from twimp import chunks <EOL> from twimp . error import ProtocolContractError , UnexpectedStatusError <EOL> from twimp . error import CommandResultError <EOL> from twimp . error import CallResultError , CallAbortedException <EOL> from twimp . proto import DispatchProtocol , DispatchFactory <EOL> from twimp . utils import ms_time_wrapped <EOL> LOG_CATEGORY = '<STR_LIT>' <EOL> import twimp . log <EOL> log = twimp . log . get_logger ( LOG_CATEGORY ) <EOL> class CancellableCallQueue ( object ) : <EOL> def __init__ ( self , reactor = reactor ) : <EOL> self . reactor = reactor <EOL> self . pending = { } <EOL> self . _next_key = <NUM_LIT:0> <EOL> def callLater ( self , delay , f , * args , ** kw ) : <EOL> call_key , self . _next_key = self . _next_key , self . _next_key + <NUM_LIT:1> <EOL> clid = self . reactor . callLater ( delay , self . _call_wrapper , call_key , f , <EOL> args , kw ) <EOL> self . pending [ call_key ] = clid <EOL> return call_key , clid <EOL> def cancel ( self , ( call_key , clid ) ) : <EOL> self . pending . pop ( call_key , None ) <EOL> return clid . cancel ( ) <EOL> def _call_wrapper ( self , call_key , f , args , kw ) : <EOL> del self . pending [ call_key ] <EOL> f ( * args , ** kw ) <EOL> def cancel_all ( self ) : <EOL> remaining = self . pending . copy ( ) <EOL> self . pending . clear ( ) <EOL> for key , clid in remaining . iteritems ( ) : <EOL> if clid . active ( ) : <EOL> clid . cancel ( ) <EOL> class DeferredTracker ( object ) : <EOL> init_trans_id = <NUM_LIT:1> <EOL> def __init__ ( self ) : <EOL> self . _pending = { } <EOL> self . _next_trans_id = { } <EOL> def next_trans ( self , key ) : <EOL> trans_id = self . _next_trans_id . setdefault ( key , self . init_trans_id ) <EOL> r , trans_id = trans_id , trans_id + <NUM_LIT:1> <EOL> self . _next_trans_id [ key ] = trans_id <EOL> return r <EOL> def push_deferred ( self , key , trans_id , d ) : <EOL> if key not in self . _pending : <EOL> self . _pending [ key ] = { } <EOL> key_queue = self . _pending [ key ] <EOL> key_queue [ trans_id ] = d <EOL> def pop_deferred ( self , key , trans_id ) : <EOL> d = None <EOL> key_queue = self . _pending . get ( key , None ) <EOL> if key_queue : <EOL> d = key_queue . pop ( trans_id , None ) <EOL> return d <EOL> def iter_all ( self ) : <EOL> return ( ( key , d ) <EOL> for ( key , key_queue ) in self . _pending . iteritems ( ) <EOL> for d in key_queue . itervalues ( ) ) <EOL> def clear ( self ) : <EOL> self . _pending = { } <EOL> def reset ( self ) : <EOL> self . _next_trans_id = { } <EOL> class StatusEventTracker ( object ) : <EOL> def __init__ ( self ) : <EOL> self . _event_callbacks = { } <EOL> def add ( self , key , code , d ) : <EOL> q = self . _event_callbacks . get ( key , None ) <EOL> if q is None : <EOL> q = self . _event_callbacks [ key ] = deque ( ) <EOL> q . append ( ( code , d ) ) <EOL> def pop ( self , key ) : <EOL> code , d = None , None <EOL> q = self . _event_callbacks . get ( key , None ) <EOL> if q : <EOL> code , d = q . popleft ( ) <EOL> return code , d <EOL> def pop_all ( self ) : <EOL> ret = ( ( code , d ) <EOL> for q in self . _event_callbacks . itervalues ( ) <EOL> for ( code , d ) in q ) <EOL> self . _event_callbacks = { } <EOL> return ret <EOL> def wait ( self , key , code ) : <EOL> d = defer . Deferred ( ) <EOL> self . add ( key , code , d ) <EOL> return d <EOL> def cancel_all ( self , reason = None ) : <EOL> waiting = self . pop_all ( ) <EOL> if reason : <EOL> for _code , d in waiting : <EOL> d . errback ( reason ) <EOL> def dispatch ( self , key , info , miss_h = None ) : <EOL> code , d = self . pop ( key ) <EOL> if not d : <EOL> if miss_h : <EOL> miss_h ( ) <EOL> else : <EOL> try : <EOL> evt_code = info . code <EOL> except AttributeError , e : <EOL> d . errback ( ProtocolContractError ( e ) ) <EOL> else : <EOL> if code is None or evt_code == code : <EOL> d . callback ( info ) <EOL> else : <EOL> d . errback ( UnexpectedStatusError ( info ) ) <EOL> class CommandDispatchProtocol ( DispatchProtocol ) : <EOL> def __init__ ( self ) : <EOL> DispatchProtocol . __init__ ( self ) <EOL> self . _cc_queue = CancellableCallQueue ( ) <EOL> self . _call_tracker = DeferredTracker ( ) <EOL> def doCommand ( self , ts , ms_id , args ) : <EOL> cmd = args [ <NUM_LIT:0> ] <EOL> handler_m = getattr ( self , '<STR_LIT>' % ( cmd , ) , None ) <EOL> if handler_m is None : <EOL> self . _cc_queue . callLater ( <NUM_LIT:0> , self . unknownCommandType , cmd , ts , <EOL> ms_id , args [ <NUM_LIT:1> : ] ) <EOL> else : <EOL> self . _cc_queue . callLater ( <NUM_LIT:0> , self . _handler_wrapper , handler_m , <EOL> ts , ms_id , args [ <NUM_LIT:1> : ] ) <EOL> def _handler_wrapper ( self , handler , ts , ms_id , args ) : <EOL> handler ( ts , ms_id , * args ) <EOL> def command__result ( self , ts , ms_id , trans_id , * args ) : <EOL> d = self . _call_tracker . pop_deferred ( ms_id , trans_id ) <EOL> if d : <EOL> d . callback ( args ) <EOL> else : <EOL> self . unexpectedCallResult ( ts , ms_id , trans_id , args ) <EOL> def command__error ( self , ts , ms_id , trans_id , * args ) : <EOL> d = self . _call_tracker . pop_deferred ( ms_id , trans_id ) <EOL> if d : <EOL> d . errback ( failure . Failure ( CommandResultError ( * args ) ) ) <EOL> else : <EOL> self . unexpectedCallError ( ts , ms_id , trans_id , args ) <EOL> def unknownCommandType ( self , cmd , ts , msid , args ) : <EOL> raise NotImplementedError ( '<STR_LIT>' % ( cmd , <EOL> ( ts , msid , args ) ) ) <EOL> def unexpectedCallResult ( self , ts , ms_id , trans_id , args ) : <EOL> log . warning ( '<STR_LIT>' , <EOL> ts , ms_id , trans_id , args ) <EOL> def unexpectedCallError ( self , ts , ms_id , trans_id , args ) : <EOL> log . warning ( '<STR_LIT>' , <EOL> ts , ms_id , trans_id , args ) <EOL> def connectionLost ( self , reason = protocol . connectionDone ) : <EOL> self . _cc_queue . cancel_all ( ) <EOL> pending_calls = list ( self . _call_tracker . iter_all ( ) ) <EOL> self . _call_tracker . clear ( ) <EOL> for ms_id , d in pending_calls : <EOL> d . errback ( reason ) <EOL> DispatchProtocol . connectionLost ( self , reason ) <EOL> def encode_amf ( self , * args ) : <EOL> return amf0 . encode ( * args ) <EOL> def _send_command ( self , ts , ms_id , body , track_id ) : <EOL> d = None <EOL> if track_id : <EOL> d = defer . Deferred ( ) <EOL> self . _call_tracker . push_deferred ( ms_id , track_id , d ) <EOL> self . muxer . sendMessage ( <NUM_LIT:0> , chunks . MSG_COMMAND , ms_id , body ) <EOL> return d <EOL> def _sendRemote ( self , ms_id , cmd , args , kwargs , track ) : <EOL> trans_id = <NUM_LIT:0> <EOL> if track : <EOL> trans_id = self . _call_tracker . next_trans ( ms_id ) <EOL> encoded_args = self . encode_amf ( cmd , trans_id , * args ) <EOL> return self . _send_command ( <NUM_LIT:0> , ms_id , encoded_args , trans_id ) <EOL> def callRemote ( self , ms_id , cmd , * args , ** kw ) : <EOL> return self . _sendRemote ( ms_id , cmd , args , kw , True ) <EOL> def signalRemote ( self , ms_id , cmd , * args , ** kw ) : <EOL> return self . _sendRemote ( ms_id , cmd , args , kw , False ) <EOL> class CommandDispatchFactory ( DispatchFactory ) : <EOL> protocol = CommandDispatchProtocol <EOL> class EventDispatchProtocol ( CommandDispatchProtocol ) : <EOL> def __init__ ( self ) : <EOL> CommandDispatchProtocol . __init__ ( self ) <EOL> self . _events = StatusEventTracker ( ) <EOL> def _onStatus_ev_key ( self , ms_id ) : <EOL> return ( None , ms_id ) <EOL> def waitStatus ( self , ms_id , code ) : <EOL> return self . _events . wait ( self . _onStatus_ev_key ( ms_id ) , code ) <EOL> def command_onStatus ( self , ts , ms_id , _trans_id , _none , info ) : <EOL> def miss_handler ( ) : <EOL> self . unhandledOnStatus ( ts , ms_id , info ) <EOL> self . _events . dispatch ( self . _onStatus_ev_key ( ms_id ) , info , <EOL> miss_h = miss_handler ) <EOL> def unhandledOnStatus ( self , ts , ms_id , info ) : <EOL> log . warning ( '<STR_LIT>' , <EOL> ts , ms_id , info ) <EOL> def connectionLost ( self , reason = protocol . connectionDone ) : <EOL> self . _events . cancel_all ( reason = reason ) <EOL> CommandDispatchProtocol . connectionLost ( self , reason ) <EOL> class EventDispatchFactory ( CommandDispatchFactory ) : <EOL> protocol = EventDispatchProtocol <EOL> class CallDispatchProtocol ( EventDispatchProtocol ) : <EOL> def __init__ ( self ) : <EOL> EventDispatchProtocol . __init__ ( self ) <EOL> def session_time ( self ) : <EOL> return time . time ( ) - self . session_init_time <EOL> def unknownCommandType ( self , cmd , ts , ms_id , args ) : <EOL> trans_id = args [ <NUM_LIT:0> ] <EOL> handler_m = getattr ( self , '<STR_LIT>' % ( cmd , ) , None ) <EOL> if handler_m is None : <EOL> d = defer . maybeDeferred ( self . unknownRemoteCall , cmd , ts , ms_id , <EOL> args [ <NUM_LIT:1> : ] ) <EOL> else : <EOL> d = defer . maybeDeferred ( handler_m , ts , ms_id , * args [ <NUM_LIT:1> : ] ) <EOL> if trans_id : <EOL> d . addCallback ( self . _remote_handler_cb , ms_id , trans_id ) <EOL> d . addErrback ( self . _remote_abort_handler_eb ) <EOL> d . addErrback ( self . _remote_handler_eb , ms_id , trans_id ) <EOL> def _remote_abort_handler_eb ( self , failure ) : <EOL> failure . trap ( CallAbortedException ) <EOL> log . debug ( '<STR_LIT>' , failure . value ) <EOL> def _remote_handler_cb ( self , result , ms_id , trans_id ) : <EOL> if not isinstance ( result , ( tuple , list ) ) : <EOL> result = ( result , ) <EOL> body = self . encode_amf ( '<STR_LIT>' , trans_id , * result ) <EOL> ts = ms_time_wrapped ( self . session_time ( ) ) <EOL> self . muxer . sendMessage ( ts , chunks . MSG_COMMAND , ms_id , body ) <EOL> def _remote_handler_eb ( self , failure , ms_id , trans_id ) : <EOL> if log . isEnabledFor ( logging . DEBUG ) : <EOL> log . info ( '<STR_LIT>' , failure . value , <EOL> exc_info = ( failure . type , failure . value , <EOL> failure . getTracebackObject ( ) ) ) <EOL> else : <EOL> log . info ( '<STR_LIT>' , failure . value ) <EOL> fatal = False <EOL> if failure . check ( CallResultError ) : <EOL> body = self . encode_amf ( '<STR_LIT>' , trans_id , <EOL> * failure . value . get_error_args ( ) ) <EOL> fatal = failure . value . is_fatal <EOL> else : <EOL> err = amf0 . Object ( code = '<STR_LIT>' , level = '<STR_LIT:error>' , <EOL> description = repr ( failure . value ) ) <EOL> body = self . encode_amf ( '<STR_LIT>' , trans_id , None , err ) <EOL> ts = ms_time_wrapped ( self . session_time ( ) ) <EOL> self . muxer . sendMessage ( ts , chunks . MSG_COMMAND , ms_id , body ) <EOL> if fatal : <EOL> self . transport . loseConnection ( ) <EOL> def unknownRemoteCall ( self , cmd , ts , ms_id , args ) : <EOL> log . warning ( '<STR_LIT>' , cmd , args ) <EOL> raise CallAbortedException ( '<STR_LIT>' % ( cmd , ) ) <EOL> class CallDispatchFactory ( EventDispatchFactory ) : <EOL> protocol = CallDispatchProtocol </s>
