twisted/protocols/stateful.py

   1 # -*- test-case-name: twisted.test.test_stateful -*-
   2
   3 # Copyright (c) Twisted Matrix Laboratories.
   4 # See LICENSE for details.
   5
   6
   7 from twisted.internet import protocol
   8
   9 try:
  10     from cStringIO import StringIO
  11 except ImportError:
  12     from StringIO import StringIO
  13
  14 class StatefulProtocol(protocol.Protocol):
  15     """A Protocol that stores state for you.
  16
  17     state is a pair (function, num_bytes). When num_bytes bytes of data arrives
  18     from the network, function is called. It is expected to return the next
  19     state or None to keep same state. Initial state is returned by
  20     getInitialState (override it).
  21     """
  22     _sful_data = None, None, 0
  23
  24     def makeConnection(self, transport):
  25         protocol.Protocol.makeConnection(self, transport)
  26         self._sful_data = self.getInitialState(), StringIO(), 0
  27
  28     def getInitialState(self):
  29         raise NotImplementedError
  30
  31     def dataReceived(self, data):
  32         state, buffer, offset = self._sful_data
  33         buffer.seek(0, 2)
  34         buffer.write(data)
  35         blen = buffer.tell() # how many bytes total is in the buffer
  36         buffer.seek(offset)
  37         while blen - offset >= state[1]:
  38             d = buffer.read(state[1])
  39             offset += state[1]
  40             next = state[0](d)
  41             if self.transport.disconnecting: # XXX: argh stupid hack borrowed right from LineReceiver
  42                 return # dataReceived won't be called again, so who cares about consistent state
  43             if next:
  44                 state = next
  45         if offset != 0:
  46             b = buffer.read()
  47             buffer.seek(0)
  48             buffer.truncate()
  49             buffer.write(b)
  50             offset = 0
  51         self._sful_data = state, buffer, offset
  52