cleanup

Signed-off-by: Mandana Vaziri <[email protected]>
IBM · Sep 9, 2024 · 21cb8be · 21cb8be
1 parent 3fca3f0
commit 21cb8be
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 4 deletions.
diff --git 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
@@ -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
@@ -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, "")