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(): def get_color_map():
if platform.system().lower() != "windows": if platform.system().lower() != "windows":
color_map = { color_map = {
"success": Fore.GREEN, "success": "green",
"failure": Fore.RED, "failure": "red",
"status": Fore.LIGHTGREEN_EX, "status": "light_green",
"code": Fore.LIGHTBLUE_EX, "code": "light_blue",
"warning": Fore.YELLOW, "warning": "yellow",
"output": Fore.LIGHTCYAN_EX, "output": "cyan",
"info": Fore.CYAN "info": "cyan"
} }
else: else:
color_map = { color_map = {
@ -54,21 +54,28 @@ def pretty_print(text, color="info"):
if color not in color_map: if color not in color_map:
color = "info" color = "info"
print(colored(text, color_map[color])) 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. 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. It use a daemon thread to run the animation. This will not block the main thread.
Color are the same as pretty_print. Color are the same as pretty_print.
""" """
global thinking_toggle global thinking_toggle
thinking_toggle = False thinking_toggle = True
def _animate(): def _animate():
global thinking_toggle global thinking_toggle
color_map = get_color_map() color_map = {
fore_color, term_color = color_map.get(color, color_map["info"]) "success": (Fore.GREEN, "green"),
#spinner = itertools.cycle(['⠋', '⠙', '⠹', '⠸', '⠼', '⠴', '⠦', '⠧', '⠇', '⠏']) "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([ spinner = itertools.cycle([
'█▒▒▒▒▒', '██▒▒▒▒', '███▒▒▒', '████▒▒', '█████▒', '██████', '█▒▒▒▒▒', '██▒▒▒▒', '███▒▒▒', '████▒▒', '█████▒', '██████',
'█████▒', '████▒▒', '███▒▒▒', '██▒▒▒▒', '█▒▒▒▒▒', '▒█▒▒▒█', '█████▒', '████▒▒', '███▒▒▒', '██▒▒▒▒', '█▒▒▒▒▒', '▒█▒▒▒█',
@ -87,7 +94,6 @@ def animate_thinking(text, color="status", duration=30):
print(colored(f"{symbol} {text}", term_color), flush=True) print(colored(f"{symbol} {text}", term_color), flush=True)
time.sleep(0.1) time.sleep(0.1)
print("\033[1A\033[K", end="", flush=True) print("\033[1A\033[K", end="", flush=True)
thinking_toggle = True
animation_thread = threading.Thread(target=_animate, daemon=True) animation_thread = threading.Thread(target=_animate, daemon=True)
animation_thread.start() animation_thread.start()
@ -106,4 +112,17 @@ def timer_decorator(func):
end_time = time() end_time = time()
pretty_print(f"{func.__name__} took {end_time - start_time:.2f} seconds to execute", "status") pretty_print(f"{func.__name__} took {end_time - start_time:.2f} seconds to execute", "status")
return result return result
return wrapper 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")