<commit_before>import zmq
from errno import EAGAIN
from collections import deque
from diesel.util.queue import Queue

zctx = zmq.Context()

class DieselZMQSocket(object):
    '''Integrate zmq's super fast event loop with ours.
    '''
    def __init__(self, socket, bind=None, connect=None):
        self.socket = socket
        from diesel.runtime import current_app
        from diesel.hub import IntWrap

        if bind:
            assert not connect
            self.socket.bind(bind)
        elif connect:
            assert not bind
            self.socket.connect(connect)

        self.hub = current_app.hub
        self.fd = IntWrap(self.socket.getsockopt(zmq.FD))
        self.outgoing = deque()
        self.incoming = Queue()

        self.hub.register(self.fd, self.handle_ready, self.flush_pending, self.error)

    def send(self, message):
        if self.outgoing:
            self.outgoing.appendleft(message)
        else:
            try:
                self.socket.send(message, zmq.NOBLOCK)
            except zmq.ZMQError, e:
                if e.errno == EAGAIN:
                    self.outgoing.appendleft(message)
                    self.hub.enable_write(self.fd)
                else:
                    raise

    def recv(self):
        return self.incoming.get()

    def handle_ready(self):
        while True:
            try:
                msg = self.socket.recv(zmq.NOBLOCK)
            except zmq.ZMQError, e:
                if e.errno == EAGAIN:
                    return
                else:
                    raise
            else:
                self.incoming.put(msg)

    def error(self):
        raise RuntimeError("OH NOES, some weird zeromq FD error callback")

    def flush_pending(self):
        while self.outgoing:
            i = self.outgoing.pop()
            try:
                self.socket.send(i, zmq.NOBLOCK)
            except zmq.ZMQError, e:
                self.outgoing.append(i)
                if e.errno == EAGAIN:
                    return
                else:
                    raise

        if not self.outgoing:
            self.hub.disable_write(self.fd)

    def __del__(self):
        self.hub.unregister(self.fd)<commit_msg>Updated ZMQ to use edge-triggered fd correctly, and to have context manager for closing.<commit_after>import zmq
from errno import EAGAIN
from collections import deque
from diesel.util.queue import Queue
from diesel.util.event import Event

zctx = zmq.Context()

class DieselZMQSocket(object):
    '''Integrate zmq's super fast event loop with ours.
    '''
    def __init__(self, socket, bind=None, connect=None, linger_time=1000):
        self.socket = socket
        from diesel.runtime import current_app
        from diesel.hub import IntWrap

        if bind:
            assert not connect
            self.socket.bind(bind)
        elif connect:
            assert not bind
            self.socket.connect(connect)

        self.hub = current_app.hub
        self.fd = IntWrap(self.socket.getsockopt(zmq.FD))

        self.write_gate = Event()
        self.read_gate = Event()
        self.linger_time = linger_time
        self.hub.register(self.fd, self.handle_transition, self.error, self.error)
        self.handle_transition()
        self.destroyed = False

    def send(self, message):
        while True:
            self.write_gate.wait()
            try:
                self.socket.send(message, zmq.NOBLOCK)
            except zmq.ZMQError, e:
                if e.errno == EAGAIN:
                    self.handle_transition() # force re-evaluation of EVENTS
                else:
                    raise
            else:
                break

    def recv(self):
        while True:
            self.read_gate.wait()
            try:
                m = self.socket.recv(zmq.NOBLOCK)
            except zmq.ZMQError, e:
                if e.errno == EAGAIN:
                    self.handle_transition() # force re-evaluation of EVENTS
                else:
                    raise
            else:
                return m

    def handle_transition(self):
        '''Handle state change.
        '''

        events = self.socket.getsockopt(zmq.EVENTS)
        if events & zmq.POLLIN:
            self.read_gate.set()
        else:
            self.read_gate.clear()

        if events & zmq.POLLOUT:
            self.write_gate.set()
        else:
            self.write_gate.clear()

    def error(self):
        raise RuntimeError("OH NOES, some weird zeromq FD error callback")

    def __del__(self):
        if not self.destroyed:
            self.hub.unregister(self.fd)
            self.socket.close(self.linger_time)

    def __enter__(self):
        return self

    def __exit__(self, *args):
        if not self.destroyed:
            self.hub.unregister(self.fd)
            self.socket.close(self.linger_time)
{'fraction_non_alphanumeric': 0.08627773212818406, 'fraction_numerical': 0.0008216926869350862, 'mean_word_length': 1.5816542948038177, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '3277483', 'n_tokens_mistral': 1421, 'n_tokens_neox': 1340, 'n_words': 328}