main: improve REPL use (still broken though).

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()
 
+    input_stream = REPLInput()
+
+    if args.mock:
         print("Source:\n", data)
-    tokens: collections.abc.Iterator[Token] = \
-        Lexer(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: