<commit_before>from flask_sockets import Sockets
from flask import Flask
from redis import StrictRedis
from json import loads, dumps
from multiprocessing import Process
from gevent import sleep, Greenlet
from geventwebsocket.exceptions import WebSocketError

processes = {}

def socket_sentinel(ws, ps):
    for msg in ps.listen():
        print msg
        if msg and isinstance(msg.get('data'), (str, unicode)):
            ws.send(msg['data'])
        
def enable_chutes(app):
    '''
    Factory method to add the chutes socket endpoint to your existing Flask app
    
    Input:
        app - Flask App Object to be extended
        
    Returns:
        None
    '''
    
    assert isinstance(app, Flask)
    
    connection = app.config['REDIS_CONN']
    r = StrictRedis(**connection)
    sockets = Sockets(app)
    
    @sockets.route('/chutes')
    def _chutes(ws):
        try:
            
            i = 0
            redis_key = None
            channel = None
    
            while True:
                
                if i == 0:
                    msg = ws.receive()
                    print msg
                    sign_on = loads(msg)
                    channel = sign_on['channel']
                    if channel not in processes:
                        processes[channel] = []
                        
                    
                    redis_key = 'c:%s'%channel
                    i += 1
                    
                    ps = r.pubsub()
                    ps.subscribe(redis_key)
                    process = Greenlet(socket_sentinel, *(ws, ps))
                    process.start()
                    processes[channel].append(process)
                    
                resp = r.blpop(redis_key, 30)
                print resp
                
                if ws.closed:
                    print 'Websocket Connection Closed by Client'
                    break
                
                if resp and isinstance(resp[-1], (str, unicode)):
                    print 'WS:', channel, '->', resp[-1]
                    ws.send(resp[-1])
                else:
                    ws.send(dumps({'data':None}))
        except WebSocketError, e:
            print dir(e)
            _processes = processes[channel]
            
            for process in _processes:
                process.kill()

class Chute(object):
    def __init__(self, channel, **kwargs):
        self.r = StrictRedis(**kwargs)
        self.channel = channel
        self._r_key = 'c:%s'%channel
        
    
    def send(self, data, timeout=90):
        self.r.lpush(self._r_key, dumps({'data':data}))
        self.r.expire(self._r_key, timeout)
        
    def publish(self, data):
        self.r.publish(self._r_key, dumps({'data':data}))
    
            

def send_response_to_chute(channel, data, **kwargs):
    r = StrictRedis(**kwargs)
    r_key = 'c:%s'%channel
    r.lpush(r_key, dumps({'data':data}))
    r.expire(r_key, kwargs.pop('timeout', 90))
    <commit_msg>Cleanup<commit_after>from flask_sockets import Sockets
from flask import Flask
from redis import StrictRedis
from json import loads, dumps
from multiprocessing import Process
from gevent import sleep, Greenlet
from geventwebsocket.exceptions import WebSocketError

processes = {}

def enable_chutes(app):
    '''
    Factory method to add the chutes socket endpoint to your existing Flask app
    
    Input:
        app - Flask App Object to be extended
        
    Returns:
        None
    '''
    
    assert isinstance(app, Flask)
    
    connection = app.config['REDIS_CONN']
    r = StrictRedis(**connection)
    sockets = Sockets(app)
    
    @sockets.route('/chutes')
    def _chutes(ws):
        try:
            
            i = 0
            redis_key = None
            channel = None
            
            while True:
                
                if i == 0:
                    msg = ws.receive()
                    print msg
                    sign_on = loads(msg)
                    channel = sign_on['channel']
                    if channel not in processes:
                        processes[channel] = []

                    redis_key = 'c:%s'%channel
                    i += 1
                    
                    ps = r.pubsub()
                    ps.subscribe(redis_key)
                    process = Greenlet(socket_sentinel, *(ws, ps))
                    process.start()
                    processes[channel].append(process)
                    
                resp = r.blpop(redis_key, 30)
                print resp
                
                if ws.closed:
                    print 'Websocket Connection Closed by Client'
                    break
                
                if resp and isinstance(resp[-1], (str, unicode)):
                    print 'WS:', channel, '->', resp[-1]
                    ws.send(resp[-1])
                else:
                    ws.send(dumps({'data':None}))
                    
        except WebSocketError, e:
            _processes = processes[channel]            
            for process in _processes:
                process.kill()

class Chute(object):
    def __init__(self, channel, **kwargs):
        self.r = StrictRedis(**kwargs)
        self.channel = channel
        self._r_key = 'c:%s'%channel
        
    def send(self, data, timeout=90):
        self.r.lpush(self._r_key, dumps({'data':data}))
        self.r.expire(self._r_key, timeout)
        
    def publish(self, data):
        self.r.publish(self._r_key, dumps({'data':data}))
    
    
def send_response_to_chute(channel, data, **kwargs):
    r = StrictRedis(**kwargs)
    r_key = 'c:%s'%channel
    r.lpush(r_key, dumps({'data':data}))
    r.expire(r_key, kwargs.pop('timeout', 90))
    

def socket_sentinel(ws, ps):
    for msg in ps.listen():
        print msg
        if msg and isinstance(msg.get('data'), (str, unicode)):
            ws.send(msg['data'])
        
{'fraction_non_alphanumeric': 0.09981515711645102, 'fraction_numerical': 0.004032935641068728, 'mean_word_length': 1.3894018466479325, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4086669', 'n_tokens_mistral': 1714, 'n_tokens_neox': 1615, 'n_words': 398}