From 974095e2cd67ccbc99a721774c6f3c2be539245f Mon Sep 17 00:00:00 2001
From: Louis Mandel <lmandel@us.ibm.com>
Date: Fri, 6 Jun 2025 12:04:43 -0400
Subject: [PATCH] chore: remove dependency on termcolor types

Signed-off-by: Louis Mandel <lmandel@us.ibm.com>
---
 src/pdl/pdl_scheduler.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/src/pdl/pdl_scheduler.py b/src/pdl/pdl_scheduler.py
index 53a795042..2ce0d0d56 100644
--- a/src/pdl/pdl_scheduler.py
+++ b/src/pdl/pdl_scheduler.py
@@ -1,14 +1,13 @@
 from typing import Any, Optional
 
 from termcolor import colored
-from termcolor._types import Color
 
 from .pdl_ast import BlockKind
 from .pdl_utils import stringify
 
 
 def color_of(kind: BlockKind):
-    color: Optional[Color]
+    color: Optional[str]
     match kind:
         case BlockKind.FUNCTION:
             color = None
@@ -52,7 +51,7 @@ def color_of(kind: BlockKind):
 
 
 def color_of_role(role: str):
-    color: Optional[Color] = None
+    color: Optional[str] = None
     match role:
         case "assistant":
             color = "green"