SO BLOODY CLOSE

lib/nntp/tiny/buffer.py

@@ -5,66 +5,64 @@ class BufferOverflow(Exception):
         super().__init__("Buffer overflow")
 
 class LineBuffer():
-    __slots__ = 'buf', 'offset',
+    __slots__ = 'buf', 'size', 'offset', 'count', 'eof', 'done',
 
     BUFFER_SIZE = 4096
 
     def __init__(self, size: int=BUFFER_SIZE):
         self.buf:    bytearray = bytearray(size)
+        self.size:   int = size
         self.offset: int = 0
+        self.count:  int = 0
+        self.eof:    bool = False
+        self.done:   bool = False
 
-    def readline(self, sock: socket.socket) -> bytes:
+    def _shift(self):
+        count = self.size - self.offset
+
+        self.buf[0:count] = self.buf[self.offset:self.count]
+
+        self.offset = 0
+        self.count  = count
+
+    def _is_full(self):
+        return self.count == self.size
+
+    def _fill(self, sock: socket.socket):
+        if self.eof:
+            return
+
+        readlen = self.size - self.count
+        data    = sock.recv(readlen)
+
+        if data == b'':
+            self.eof = True
+            return
+
+        readlen = len(data)
+
+        self.buf[self.offset:self.offset+readlen] = data
+
+        self.count += readlen
+
+    def _flush(self, end: int):
+        ret = self.buf[self.offset:end+1]
+
+        self.offset = end + 1
+
+        return str(ret, 'ascii')
+
+    def readline(self, sock: socket.socket) -> str:
         while True:
-            #
-            # Check and see if there is already a line in the
-            # internal buffer.
-            #
-            offset = self.buf.find(b'\n', self.offset, len(self.buf))
+            index = self.buf.find(b'\n', self.offset, self.count)
 
-            if offset < 0:
-                #
-                # There is no line in the buffer.  Check and see if
-                # there is more room in the buffer to continue
-                # reading.
-                #
-                if offset == self.BUFFER_SIZE - 1:
-                    #
-                    # There is no room left in the buffer; the line
-                    # must be too long.  Panic.
-                    #
-                    raise BufferOverflow()
-
-                #
-                # Read from the socket to attempt to fill the buffer
-                # more.
-                #
-                data = sock.recv(self.BUFFER_SIZE - self.offset)
-
-                #
-                # Check to see if recv() returned any data.  If not,
-                # the socket is closed.  Return an empty string.
-                #
-                if data == b'':
-                    return ''
-
-                count = len(data)
-
-                #
-                # Now, fill the buffer.  The next iteration of the
-                # loop will be used to check for a newline.
-                #
-                self.buf[self.offset:self.offset+count] = data
+            if index < 0:
+                if self._is_full():
+                    if self.offset == 0:
+                        raise BufferOverflow()
+                    else:
+                        self._shift()
+                else:
+                    self._fill(sock)
             else:
-                #
-                # There is a line in the buffer.  Prepare to return
-                # it.
-                #
-                ret = self.buf[self.offset:offset]
-
-                #
-                # Advance the offset to one character beyond the
-                # newline.
-                #
-                self.offset = offset + 1
-
-                return str(ret, 'ascii')
+                return self._flush(index)