transport code clean up

thriftpy/transport/__init__.py

@@ -25,7 +25,7 @@ def readall(read_fn, sz):
 class TTransportBase(object):
     """Base class for Thrift transport layer."""
 
-    def _read(self):
+    def _read(self, sz):
         raise NotImplementedError
 
     def read(self, sz):

thriftpy/transport/buffered/__init__.py

@@ -17,40 +17,40 @@ class TBufferedTransport(TTransportBase):
     DEFAULT_BUFFER = 4096
 
     def __init__(self, trans, buf_size=DEFAULT_BUFFER):
-        self.__trans = trans
-        self.__wbuf = BytesIO()
-        self.__rbuf = BytesIO(b"")
-        self.__buf_size = buf_size
+        self._trans = trans
+        self._wbuf = BytesIO()
+        self._rbuf = BytesIO(b"")
+        self._buf_size = buf_size
 
     def is_open(self):
-        return self.__trans.is_open()
+        return self._trans.is_open()
 
     def open(self):
-        return self.__trans.open()
+        return self._trans.open()
 
     def close(self):
-        return self.__trans.close()
+        return self._trans.close()
 
     def _read(self, sz):
-        ret = self.__rbuf.read(sz)
+        ret = self._rbuf.read(sz)
         if len(ret) != 0:
             return ret
 
-        self.__rbuf = BytesIO(self.__trans.read(max(sz, self.__buf_size)))
-        return self.__rbuf.read(sz)
+        self._rbuf = BytesIO(self._trans.read(max(sz, self._buf_size)))
+        return self._rbuf.read(sz)
 
     def write(self, buf):
-        self.__wbuf.write(buf)
+        self._wbuf.write(buf)
 
     def flush(self):
-        out = self.__wbuf.getvalue()
+        out = self._wbuf.getvalue()
         # reset wbuf before write/flush to preserve state on underlying failure
-        self.__wbuf = BytesIO()
-        self.__trans.write(out)
-        self.__trans.flush()
+        self._wbuf = BytesIO()
+        self._trans.write(out)
+        self._trans.flush()
 
     def getvalue(self):
-        return self.__trans.getvalue()
+        return self._trans.getvalue()
 
 
 class TBufferedTransportFactory(object):

thriftpy/transport/framed/__init__.py

@@ -13,53 +13,56 @@
 class TFramedTransport(TTransportBase):
     """Class that wraps another transport and frames its I/O when writing."""
     def __init__(self, trans):
-        self.__trans = trans
-        self.__rbuf = BytesIO()
+        self._trans = trans
+        self._rbuf = BytesIO()
+        self._wbuf = BytesIO()
 
     def is_open(self):
-        return self.__trans.is_open()
+        return self._trans.is_open()
 
     def open(self):
-        return self.__trans.open()
+        return self._trans.open()
 
     def close(self):
-        return self.__trans.close()
+        return self._trans.close()
 
     def read(self, sz):
         # Important: don't attempt to read the next frame if the caller
         # doesn't actually need any data.
         if sz == 0:
             return b''
 
-        ret = self.__rbuf.read(sz)
+        ret = self._rbuf.read(sz)
         if len(ret) != 0:
             return ret
 
         self.read_frame()
-        return self.__rbuf.read(sz)
+        return self._rbuf.read(sz)
 
     def read_frame(self):
-        buff = readall(self.__trans.read, 4)
+        buff = readall(self._trans.read, 4)
         sz, = struct.unpack('!i', buff)
-        frame = readall(self.__trans.read, sz)
-        self.__rbuf = BytesIO(frame)
+        frame = readall(self._trans.read, sz)
+        self._rbuf = BytesIO(frame)
 
     def write(self, buf):
-        wsz = len(buf)
+        self._wbuf.write(buf)
+
+    def flush(self):
+        # reset wbuf before write/flush to preserve state on underlying failure
+        out = self._wbuf.getvalue()
+        self._wbuf = BytesIO()
+
         # N.B.: Doing this string concatenation is WAY cheaper than making
         # two separate calls to the underlying socket object. Socket writes in
         # Python turn out to be REALLY expensive, but it seems to do a pretty
         # good job of managing string buffer operations without excessive
         # copies
-        towrite = struct.pack("!i", wsz) + buf
-        self.__trans.write(towrite)
-        self.__trans.flush()
-
-    def flush(self):
-        pass
+        self._trans.write(struct.pack("!i", len(out)) + out)
+        self._trans.flush()
 
     def getvalue(self):
-        return self.__trans.getvalue()
+        return self._trans.getvalue()
 
 
 class TFramedTransportFactory(object):