From 3fca3f0d0e7c5a8aeec662ee936e03a3144b8530 Mon Sep 17 00:00:00 2001 From: Mandana Vaziri Date: Mon, 9 Sep 2024 12:40:40 -0400 Subject: [PATCH 1/2] coloring output Signed-off-by: Mandana Vaziri --- .gitignore | 2 ++ pdl/pdl_interpreter.py | 6 ++++-- pdl/pdl_scheduler.py | 18 ++++++++++++++++++ tests/test_line_table.py | 8 ++++++-- 4 files changed, 30 insertions(+), 4 deletions(-) diff --git a/.gitignore b/.gitignore index a9275646..88ee8825 100644 --- a/.gitignore +++ b/.gitignore @@ -147,6 +147,8 @@ pdl-live/package-lock.json # PDL trace +*_result.json +*_result.yaml *_trace.json # Built docs diff --git a/pdl/pdl_interpreter.py b/pdl/pdl_interpreter.py index 273c1ab4..bdc1cd14 100644 --- a/pdl/pdl_interpreter.py +++ b/pdl/pdl_interpreter.py @@ -63,6 +63,8 @@ YieldBackgroundMessage, YieldMessage, YieldResultMessage, + CodeYieldResultMessage, + ModelYieldResultMessage, schedule, ) from .pdl_schema_validator import type_check_args, type_check_spec @@ -286,7 +288,7 @@ def step_block_body( state, scope, block, loc ) if state.yield_result: - yield YieldResultMessage(result) + yield CodeYieldResultMessage(result) if state.yield_background: yield YieldBackgroundMessage(background) case GetBlock(get=var): @@ -889,7 +891,7 @@ def generate_client_response_streaming( role = None for chunk in msg_stream: if state.yield_result: - yield YieldResultMessage(chunk["content"]) + yield ModelYieldResultMessage(chunk["content"]) if state.yield_background: yield YieldBackgroundMessage([chunk]) if complete_msg is None: diff --git a/pdl/pdl_scheduler.py b/pdl/pdl_scheduler.py index 4fe21eb3..703f6d3f 100644 --- a/pdl/pdl_scheduler.py +++ b/pdl/pdl_scheduler.py @@ -39,6 +39,10 @@ def step_to_completion(gen: Generator[Any, Any, GeneratorReturnT]) -> GeneratorR pass return w.value +MODEL_COLOR = "\033[92m" # Green +CODE_COLOR = "\033[95m" # Purple +END_COLOR = "\033[0m" # End color + class MessageKind(Enum): RESULT = 0 @@ -53,9 +57,20 @@ class YieldMessage: @dataclass class YieldResultMessage(YieldMessage): kind = MessageKind.RESULT + color = None result: Any +@dataclass +class ModelYieldResultMessage(YieldResultMessage): + color = MODEL_COLOR + + +@dataclass +class CodeYieldResultMessage(YieldResultMessage): + color = CODE_COLOR + + @dataclass class YieldBackgroundMessage(YieldMessage): kind = MessageKind.BACKGROUND @@ -92,6 +107,9 @@ def schedule( try: msg = gen.send(v) match msg: + case ModelYieldResultMessage(result=result) | CodeYieldResultMessage(result=result): + print(msg.color + stringify(result) + END_COLOR, end="") + todo_next.append((i, gen, None)) case YieldResultMessage(result=result): print(stringify(result), end="") todo_next.append((i, gen, None)) diff --git a/tests/test_line_table.py b/tests/test_line_table.py index 8ebef126..0a1170b2 100644 --- a/tests/test_line_table.py +++ b/tests/test_line_table.py @@ -1,13 +1,17 @@ from pdl.pdl_interpreter import generate - +from pdl.pdl_scheduler import CODE_COLOR, END_COLOR, MODEL_COLOR def do_test(t, capsys): generate(t["file"], None, None, {}, None) captured = capsys.readouterr() - output = captured.out.split("\n") + output_string = remove_coloring(captured.out) + output = output_string.split("\n") print(output) assert set(output) == set(t["errors"]) +def remove_coloring(text): + return text.replace(MODEL_COLOR, "").replace(CODE_COLOR, "").replace(END_COLOR, "") + line = { "file": "tests/data/line/hello.pdl", From 21cb8be44ecf833bc275849016352dab33f985a3 Mon Sep 17 00:00:00 2001 From: Mandana Vaziri Date: Mon, 9 Sep 2024 12:47:40 -0400 Subject: [PATCH 2/2] cleanup Signed-off-by: Mandana Vaziri --- pdl/pdl_interpreter.py | 4 ++-- pdl/pdl_scheduler.py | 8 ++++++-- tests/test_line_table.py | 2 ++ 3 files changed, 10 insertions(+), 4 deletions(-) diff --git a/pdl/pdl_interpreter.py b/pdl/pdl_interpreter.py index bdc1cd14..211641e7 100644 --- a/pdl/pdl_interpreter.py +++ b/pdl/pdl_interpreter.py @@ -59,12 +59,12 @@ from .pdl_location_utils import append, get_loc_string from .pdl_parser import PDLParseError, parse_file from .pdl_scheduler import ( + CodeYieldResultMessage, ModelCallMessage, + ModelYieldResultMessage, YieldBackgroundMessage, YieldMessage, YieldResultMessage, - CodeYieldResultMessage, - ModelYieldResultMessage, schedule, ) from .pdl_schema_validator import type_check_args, type_check_spec diff --git a/pdl/pdl_scheduler.py b/pdl/pdl_scheduler.py index 703f6d3f..5d062ff7 100644 --- a/pdl/pdl_scheduler.py +++ b/pdl/pdl_scheduler.py @@ -39,9 +39,11 @@ def step_to_completion(gen: Generator[Any, Any, GeneratorReturnT]) -> GeneratorR pass return w.value + MODEL_COLOR = "\033[92m" # Green CODE_COLOR = "\033[95m" # Purple END_COLOR = "\033[0m" # End color +NO_COLOR = "" class MessageKind(Enum): @@ -57,7 +59,7 @@ class YieldMessage: @dataclass class YieldResultMessage(YieldMessage): kind = MessageKind.RESULT - color = None + color = NO_COLOR result: Any @@ -107,7 +109,9 @@ def schedule( try: msg = gen.send(v) match msg: - case ModelYieldResultMessage(result=result) | CodeYieldResultMessage(result=result): + case ModelYieldResultMessage( + result=result + ) | CodeYieldResultMessage(result=result): print(msg.color + stringify(result) + END_COLOR, end="") todo_next.append((i, gen, None)) case YieldResultMessage(result=result): diff --git a/tests/test_line_table.py b/tests/test_line_table.py index 0a1170b2..c5e28fb7 100644 --- a/tests/test_line_table.py +++ b/tests/test_line_table.py @@ -1,6 +1,7 @@ from pdl.pdl_interpreter import generate from pdl.pdl_scheduler import CODE_COLOR, END_COLOR, MODEL_COLOR + def do_test(t, capsys): generate(t["file"], None, None, {}, None) captured = capsys.readouterr() @@ -9,6 +10,7 @@ def do_test(t, capsys): print(output) assert set(output) == set(t["errors"]) + def remove_coloring(text): return text.replace(MODEL_COLOR, "").replace(CODE_COLOR, "").replace(END_COLOR, "")