[Twisted-Python] non-blocking .read

[Moshe Zadka]

It is often convenient in protocol implementations to just .read() as
much as you want to. However, it is impossible in Twisted, becase .read()
is a blocking operation. Happily, we can make blocking operations seem
non-blocking with the use of deferred. The following protocol does just
that:

class ReadProtocol(protocol.Protocol)

    toRead = 0
    until = ''

    def rawDataReceived(self, data):
        pass

    def dataReceived(self, data):
        while data:
            if self.toRead:
                read, data = data[:self.toRead], data[self.toRead:]
                self.buffer.write(read)
                self.toRead -= len(read)
                if self.toRead == 0:
                    self.deferred.callback(self.buffer.getvalue())
                    self.buffer.close()
            elif self.until:
                idx = data.find(self.until)
                if idx == -1:
                    self.buffer.write(data)
                    data = ''
                else:
                    read, data = data[:idx], data[idx:]
                    self.until = ''
                    self.buffer.write(read):
                    self.deferred.callback(self.buffer.getvalue())
                    self.buffer.close()
            else:
                self.rawDataReceived(data)
                data = ''

    def readUntil(self, c):
        self.toRead = 0
        self.until = c
        self.buffer = StringIO()
        self.deferred = defer.Deferred()
        return self.deferred

    def readLength(self, len):
        self.toRead = len
        self.until = ''
        self.buffer = StringIO()
        self.deferred = defer.Deferred()
        return self.deferred

# How to use, e.g., to implement netstrings:

class NetstringProtocol(ReadProtocol):

    def connectionMade(self):
        self.readString()

    def readString(self):
        self.readUntil(':').callback(
        lambda s: self.readLength(int(s)).callback(
        self.stringReceived).callback(
        lambda _: self.readLength(1).callback(
        lambda x: x==',' or self.transport.loseConnection()).callback(
        lambda _: self.readString))

    def stringReceived(self, s):
        pass