Add StreamLineReader with unit tests

This commit is contained in:
Romuald Juchnowicz-Bierbasz
2019-06-28 14:00:44 +02:00
parent 67e8681de6
commit 4e1ea8056d
3 changed files with 84 additions and 24 deletions

View File

@@ -5,6 +5,8 @@ import logging
import inspect
import json
from galaxy.reader import StreamLineReader
class JsonRpcError(Exception):
def __init__(self, code, message, data=None):
self.code = code
@@ -67,14 +69,12 @@ def anonymise_sensitive_params(params, sensitive_params):
class Server():
def __init__(self, reader, writer, encoder=json.JSONEncoder()):
self._active = True
self._reader = reader
self._reader = StreamLineReader(reader)
self._writer = writer
self._encoder = encoder
self._methods = {}
self._notifications = {}
self._eof_listeners = []
self._input_buffer = bytes()
self._processed_input_buffer_it = 0
def register_method(self, name, callback, internal, sensitive_params=False):
"""
@@ -106,7 +106,7 @@ class Server():
async def run(self):
while self._active:
try:
data = await self._readline()
data = await self._reader.readline()
if not data:
self._eof()
continue
@@ -117,26 +117,6 @@ class Server():
logging.debug("Received %d bytes of data", len(data))
self._handle_input(data)
async def _readline(self):
"""Like StreamReader.readline but without limit"""
while True:
# check if there is no unprocessed data in the buffer
if not self._input_buffer or self._processed_input_buffer_it != 0:
chunk = await self._reader.read(1024)
if not chunk:
return bytes() # EOF
self._input_buffer += chunk
it = self._input_buffer.find(b"\n", self._processed_input_buffer_it)
if it < 0:
self._processed_input_buffer_it = len(self._input_buffer)
continue
line = self._input_buffer[:it]
self._input_buffer = self._input_buffer[it+1:]
self._processed_input_buffer_it = 0
return line
def stop(self):
self._active = False

28
src/galaxy/reader.py Normal file
View File

@@ -0,0 +1,28 @@
from asyncio import StreamReader
class StreamLineReader:
"""Handles StreamReader readline without buffer limit"""
def __init__(self, reader: StreamReader):
self._reader = reader
self._buffer = bytes()
self._processed_buffer_it = 0
async def readline(self):
while True:
# check if there is no unprocessed data in the buffer
if not self._buffer or self._processed_buffer_it != 0:
chunk = await self._reader.read(1024)
if not chunk:
return bytes() # EOF
self._buffer += chunk
it = self._buffer.find(b"\n", self._processed_buffer_it)
if it < 0:
self._processed_buffer_it = len(self._buffer)
continue
line = self._buffer[:it]
self._buffer = self._buffer[it+1:]
self._processed_buffer_it = 0
return line