<s> import uuid <EOL> from gevent import spawn , wait <EOL> from gevent . event import Event <EOL> from gevent . monkey import patch_all <EOL> from gevent . queue import Queue , Empty <EOL> from gevent . select import select <EOL> from werkzeug . exceptions import HTTPException <EOL> from . websocket import WebSocket , WebSocketMiddleware <EOL> from . _uwsgi import uwsgi <EOL> class GeventWebSocketClient ( object ) : <EOL> def __init__ ( self , environ , fd , send_event , send_queue , recv_event , <EOL> recv_queue , timeout = <NUM_LIT:5> ) : <EOL> self . environ = environ <EOL> self . fd = fd <EOL> self . send_event = send_event <EOL> self . send_queue = send_queue <EOL> self . recv_event = recv_event <EOL> self . recv_queue = recv_queue <EOL> self . timeout = timeout <EOL> self . id = str ( uuid . uuid1 ( ) ) <EOL> self . connected = True <EOL> def send ( self , msg , binary = True ) : <EOL> if binary : <EOL> return self . send_binary ( msg ) <EOL> self . send_queue . put ( msg ) <EOL> self . send_event . set ( ) <EOL> def send_binary ( self , msg ) : <EOL> self . send_queue . put ( msg ) <EOL> self . send_event . set ( ) <EOL> def receive ( self ) : <EOL> return self . recv ( ) <EOL> def recv ( self ) : <EOL> return self . recv_queue . get ( ) <EOL> def close ( self ) : <EOL> self . connected = False <EOL> class GeventWebSocketMiddleware ( WebSocketMiddleware ) : <EOL> client = GeventWebSocketClient <EOL> def __call__ ( self , environ , start_response ) : <EOL> urls = self . websocket . url_map . bind_to_environ ( environ ) <EOL> try : <EOL> endpoint , args = urls . match ( ) <EOL> handler = self . websocket . view_functions [ endpoint ] <EOL> except HTTPException : <EOL> handler = None <EOL> if not handler or '<STR_LIT>' not in environ : <EOL> return self . wsgi_app ( environ , start_response ) <EOL> uwsgi . websocket_handshake ( environ [ '<STR_LIT>' ] , <EOL> environ . get ( '<STR_LIT>' , '<STR_LIT>' ) ) <EOL> send_event = Event ( ) <EOL> send_queue = Queue ( ) <EOL> recv_event = Event ( ) <EOL> recv_queue = Queue ( maxsize = <NUM_LIT:1> ) <EOL> client = self . client ( environ , uwsgi . connection_fd ( ) , send_event , <EOL> send_queue , recv_event , recv_queue , <EOL> self . websocket . timeout ) <EOL> handler = spawn ( handler , client , ** args ) <EOL> def listener ( client ) : <EOL> select ( [ client . fd ] , [ ] , [ ] , client . timeout ) <EOL> recv_event . set ( ) <EOL> listening = spawn ( listener , client ) <EOL> while True : <EOL> if not client . connected : <EOL> recv_queue . put ( None ) <EOL> listening . kill ( ) <EOL> handler . join ( client . timeout ) <EOL> return '<STR_LIT>' <EOL> wait ( [ handler , send_event , recv_event ] , None , <NUM_LIT:1> ) <EOL> if send_event . is_set ( ) : <EOL> try : <EOL> while True : <EOL> uwsgi . websocket_send ( send_queue . get_nowait ( ) ) <EOL> except Empty : <EOL> send_event . clear ( ) <EOL> except IOError : <EOL> client . connected = False <EOL> elif recv_event . is_set ( ) : <EOL> recv_event . clear ( ) <EOL> try : <EOL> recv_queue . put ( uwsgi . websocket_recv_nb ( ) ) <EOL> listening = spawn ( listener , client ) <EOL> except IOError : <EOL> client . connected = False <EOL> elif handler . ready ( ) : <EOL> listening . kill ( ) <EOL> return '<STR_LIT>' <EOL> class GeventWebSocket ( WebSocket ) : <EOL> middleware = GeventWebSocketMiddleware <EOL> def init_app ( self , app ) : <EOL> aggressive = app . config . get ( '<STR_LIT>' , True ) <EOL> patch_all ( aggressive = aggressive ) <EOL> super ( GeventWebSocket , self ) . init_app ( app ) </s>
