From 21cb8be44ecf833bc275849016352dab33f985a3 Mon Sep 17 00:00:00 2001
From: Mandana Vaziri <mvaziri@us.ibm.com>
Date: Mon, 9 Sep 2024 12:47:40 -0400
Subject: [PATCH] cleanup

Signed-off-by: Mandana Vaziri <mvaziri@us.ibm.com>
---
 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, "")