diff --git a/platform_utils.py b/platform_utils.py index f1da75eb..d10e888d 100644 --- a/platform_utils.py +++ b/platform_utils.py @@ -90,6 +90,11 @@ class _FileDescriptorStreamsNonBlocking(FileDescriptorStreams): """ Implementation of FileDescriptorStreams for platforms that support non blocking I/O. """ + def __init__(self): + super(_FileDescriptorStreamsNonBlocking, self).__init__() + self._poll = select.poll() + self._fd_to_stream = {} + class Stream(object): """ Encapsulates a file descriptor """ @@ -114,11 +119,18 @@ class _FileDescriptorStreamsNonBlocking(FileDescriptorStreams): self.fd.close() def _create_stream(self, fd, dest, std_name): - return self.Stream(fd, dest, std_name) + stream = self.Stream(fd, dest, std_name) + self._fd_to_stream[stream.fileno()] = stream + self._poll.register(stream, select.POLLIN) + return stream + + def remove(self, stream): + self._poll.unregister(stream) + del self._fd_to_stream[stream.fileno()] + super(_FileDescriptorStreamsNonBlocking, self).remove(stream) def select(self): - ready_streams, _, _ = select.select(self.streams, [], []) - return ready_streams + return [self._fd_to_stream[fd] for fd, _ in self._poll.poll()] class _FileDescriptorStreamsThreads(FileDescriptorStreams):