77 lines
1.8 KiB
Python
77 lines
1.8 KiB
Python
import socket
|
|
|
|
class BufferOverflow(Exception):
|
|
def __init__(self):
|
|
super().__init__("Buffer overflow")
|
|
|
|
class LineBuffer():
|
|
__slots__ = 'buf', 'size', 'offset_i', 'offset_o', 'eof', 'done',
|
|
|
|
BUFFER_SIZE = 4096
|
|
|
|
def __init__(self, size: int=BUFFER_SIZE):
|
|
self.buf: bytearray = bytearray(size)
|
|
self.size: int = size
|
|
self.offset_i: int = 0
|
|
self.offset_o: int = 0
|
|
self.eof: bool = False
|
|
self.done: bool = False
|
|
|
|
def _shift(self):
|
|
count = self.offset_i - self.offset_o
|
|
|
|
self.buf[0:count] = self.buf[self.offset_o:self.offset_i]
|
|
|
|
self.offset_i = count
|
|
self.offset_o = 0
|
|
|
|
def _is_full(self):
|
|
return self.offset_i >= self.size
|
|
|
|
def _fill(self, sock: socket.socket):
|
|
if self.eof:
|
|
return
|
|
|
|
data = sock.recv(self.size - self.offset_i)
|
|
|
|
if data == b'':
|
|
self.eof = True
|
|
return
|
|
|
|
readlen = len(data)
|
|
|
|
self.buf[self.offset_i:self.size] = data
|
|
|
|
self.offset_i += readlen
|
|
|
|
def _drain(self, end=None):
|
|
if end is None:
|
|
end = self.offset_i + 1
|
|
|
|
ret = self.buf[self.offset_o:end+1]
|
|
|
|
self.offset_o = end + 1
|
|
|
|
return str(ret, 'ascii')
|
|
|
|
def readline(self, sock: socket.socket) -> str:
|
|
if self.done:
|
|
return ''
|
|
|
|
while True:
|
|
index = self.buf.find(b'\n', self.offset_o, self.offset_i)
|
|
|
|
if index < 0:
|
|
if self._is_full():
|
|
if self.offset_o == 0:
|
|
raise BufferOverflow()
|
|
|
|
self._shift()
|
|
elif self.eof:
|
|
self.done = True
|
|
|
|
return self._drain()
|
|
else:
|
|
self._fill(sock)
|
|
else:
|
|
return self._drain(index)
|