From d2823f5abf5d7d42b56ca68895f72118b8248d2b Mon Sep 17 00:00:00 2001
From: Antoine Viallon <antoine@lesviallon.fr>
Date: Thu, 11 Jan 2024 18:50:02 +0100
Subject: [PATCH] main: improve REPL use (still broken though).

---
 compiler/__main__.py | 31 +++++++++++++++++++++++--------
 1 file changed, 23 insertions(+), 8 deletions(-)

diff --git a/compiler/__main__.py b/compiler/__main__.py
index 17f43aa..0a28bb9 100644
--- a/compiler/__main__.py
+++ b/compiler/__main__.py
@@ -1,16 +1,28 @@
 from __future__ import annotations
 
 import argparse
-import collections.abc
+import io
 import sys
+import typing
 
 from . import semantic, ir, optimizations
 from .errors import CompilationError, CompilationWarning
-from .lexer import Lexer, Tokens, Token
+from .lexer import Lexer, Tokens
 from .logger import rootLogger, LogLevel
 from .parser import Parser
 
 
+class REPLInput(io.TextIOWrapper):
+    def __init__(self, buffer: typing.IO[bytes] = sys.stdin.buffer,
+                 prompt_buffer: typing.IO[bytes] = sys.stdout.buffer):
+        super().__init__(buffer)
+        self._prompt_buffer = prompt_buffer
+
+    def readline(self, __size: int = -1) -> str:
+        self._prompt_buffer.write(b">> ")
+        return super().readline(__size)
+
+
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--mock", action="store_true", help="Use default test input")
@@ -32,12 +44,15 @@ def main():
             // { let b: uint32; }
             display(byte, 3) + 1;
             """
-    if not args.mock:
-        data = sys.stdin.read().strip()
 
-    print("Source:\n", data)
-    tokens: collections.abc.Iterator[Token] = \
-        Lexer(data,
+    input_stream = REPLInput()
+
+    if args.mock:
+        print("Source:\n", data)
+        input_stream = io.StringIO(data)
+
+    tokens = \
+        Lexer(input_stream,
               token_filter=lambda token: token.kind not in [Tokens.Blank, Tokens.Newline, Tokens.Comment])
 
     parser = Parser(tokens)
@@ -65,7 +80,7 @@ def main():
         print(ir.IRRegister.get_registers())
 
     except CompilationError as e:
-        e.location.source = data
+        e.location.source = tokens.data
         print(f"{e}\n{e.location.show_in_source()}", flush=True)
         if e.__cause__ is not None:
             if rootLogger.level <= LogLevel.Debug: