fix: animate_thinking bug

This commit is contained in:
martin legrand 2025-03-29 16:13:42 +01:00
parent 862a78276f
commit d8ded2d456

View File

@ -12,13 +12,13 @@ thinking_toggle = False
def get_color_map():
if platform.system().lower() != "windows":
color_map = {
"success": Fore.GREEN,
"failure": Fore.RED,
"status": Fore.LIGHTGREEN_EX,
"code": Fore.LIGHTBLUE_EX,
"warning": Fore.YELLOW,
"output": Fore.LIGHTCYAN_EX,
"info": Fore.CYAN
"success": "green",
"failure": "red",
"status": "light_green",
"code": "light_blue",
"warning": "yellow",
"output": "cyan",
"info": "cyan"
}
else:
color_map = {
@ -54,21 +54,28 @@ def pretty_print(text, color="info"):
if color not in color_map:
color = "info"
print(colored(text, color_map[color]))
print(' ' * 10) # prevent cut-off
def animate_thinking(text, color="status", duration=30):
def animate_thinking(text, color="status", duration=2):
"""
Animate a thinking spinner while a task is being executed.
It use a daemon thread to run the animation. This will not block the main thread.
Color are the same as pretty_print.
"""
global thinking_toggle
thinking_toggle = False
thinking_toggle = True
def _animate():
global thinking_toggle
color_map = get_color_map()
fore_color, term_color = color_map.get(color, color_map["info"])
#spinner = itertools.cycle(['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏'])
color_map = {
"success": (Fore.GREEN, "green"),
"failure": (Fore.RED, "red"),
"status": (Fore.LIGHTGREEN_EX, "light_green"),
"code": (Fore.LIGHTBLUE_EX, "light_blue"),
"warning": (Fore.YELLOW, "yellow"),
"output": (Fore.LIGHTCYAN_EX, "cyan"),
"default": (Fore.RESET, "black"),
"info": (Fore.CYAN, "cyan")
}
fore_color, term_color = color_map.get(color, color_map["default"])
spinner = itertools.cycle([
'█▒▒▒▒▒', '██▒▒▒▒', '███▒▒▒', '████▒▒', '█████▒', '██████',
'█████▒', '████▒▒', '███▒▒▒', '██▒▒▒▒', '█▒▒▒▒▒', '▒█▒▒▒█',
@ -87,7 +94,6 @@ def animate_thinking(text, color="status", duration=30):
print(colored(f"{symbol} {text}", term_color), flush=True)
time.sleep(0.1)
print("\033[1A\033[K", end="", flush=True)
thinking_toggle = True
animation_thread = threading.Thread(target=_animate, daemon=True)
animation_thread.start()
@ -107,3 +113,16 @@ def timer_decorator(func):
pretty_print(f"{func.__name__} took {end_time - start_time:.2f} seconds to execute", "status")
return result
return wrapper
if __name__ == "__main__":
import time
pretty_print("starting imaginary task", "success")
animate_thinking("Thinking...", "status")
time.sleep(4)
pretty_print("starting another task", "failure")
animate_thinking("Thinking...", "status")
time.sleep(4)
pretty_print("yet another task", "info")
animate_thinking("Thinking...", "status")
time.sleep(4)
pretty_print("This is an info message", "info")