add newline stream io

Author: aplr

Diff for: lf_toolkit/rpc/ipc_server.py

@@ -4,8 +4,11 @@
 
 import anyio
 
+from lf_toolkit.rpc.stream_io import StreamIO
+
 from .ipc_listener_base import IPCListener
 from .rpc_handler import RpcHandler
+from .stream_io import NewlineStreamIO
 from .stream_io import StreamServer
 
 
@@ -34,6 +37,9 @@ def __init__(
         self._listener = listener_factory(endpoint)
         super().__init__(handler)
 
+    def wrap_io(self, client: StreamIO) -> StreamIO:
+        return NewlineStreamIO(client)
+
     async def run(self):
         async with anyio.create_task_group() as task_group:
             async for client in self._listener.listen():

Diff for: lf_toolkit/rpc/stdio_server.py

@@ -7,6 +7,7 @@
 from anyio.streams.stapled import StapledByteStream
 
 from .rpc_handler import RpcHandler
+from .stream_io import PrefixStreamIO
 from .stream_io import StreamIO
 from .stream_io import StreamServer
 
@@ -38,5 +39,8 @@ def __init__(self, handler: Optional[RpcHandler] = None):
         super().__init__(handler)
         self._client = StdioClient()
 
+    def wrap_io(self, client: StreamIO) -> StreamIO:
+        return PrefixStreamIO(client)
+
     async def run(self):
         await self._handle_client(self._client)

Diff for: lf_toolkit/rpc/stream_io.py

@@ -17,6 +17,26 @@ async def write(self, data: bytes):
         pass
 
 
+class NewlineStreamIO:
+
+    base: StreamIO
+
+    def __init__(self, base: StreamIO):
+        self.base = base
+
+    async def read(self, size: int) -> bytes:
+        data = b""
+        while len(data) < size:
+            chunk = await self.base.read(1)
+            if chunk == b"\n":
+                break
+            data += chunk
+        return data
+
+    async def write(self, data: bytes):
+        await self.base.write(data + b"\n")
+
+
 class PrefixStreamIO:
 
     base: StreamIO
@@ -56,12 +76,15 @@ async def write(self, data: bytes):
 
 class StreamServer(BaseServer):
 
+    def wrap_io(self, client: StreamIO) -> StreamIO:
+        return client
+
     async def _handle_client(self, client: StreamIO):
-        io = PrefixStreamIO(client)
+        io = self.wrap_io(client)
 
         while True:
             try:
-                data = await io.read(1024)
+                data = await io.read(4096)
 
                 if not data:
                     print("Received empty data")