feat: Implement terminal cleanup on exit and signal handling in ComfyConsole

This commit is contained in:
Andy
2025-07-31 18:25:18 +00:00
parent 8c798b95c4
commit 091d7335a3
2 changed files with 43 additions and 2 deletions

View File

@@ -765,7 +765,8 @@ class dl:
DOWNLOAD_LICENCE_ONLY.set()
try:
with Live(Padding(download_table, (1, 5)), console=console, refresh_per_second=5):
# Use transient mode to prevent display remnants
with Live(Padding(download_table, (1, 5)), console=console, refresh_per_second=5, transient=True):
with ThreadPoolExecutor(downloads) as pool:
for download in futures.as_completed(
(
@@ -1009,7 +1010,7 @@ class dl:
multiplex_tasks.append((task_id, task_tracks))
with Live(Padding(progress, (0, 5, 1, 5)), console=console):
with Live(Padding(progress, (0, 5, 1, 5)), console=console, transient=True):
for task_id, task_tracks in multiplex_tasks:
progress.start_task(task_id) # TODO: Needed?
muxed_path, return_code, errors = task_tracks.mux(

View File

@@ -1,4 +1,7 @@
import atexit
import logging
import signal
import sys
from datetime import datetime
from types import ModuleType
from typing import IO, Callable, Iterable, List, Literal, Mapping, Optional, Union
@@ -167,6 +170,8 @@ class ComfyConsole(Console):
time.monotonic.
"""
_cleanup_registered = False
def __init__(
self,
*,
@@ -233,6 +238,9 @@ class ComfyConsole(Console):
if log_renderer:
self._log_render = log_renderer
# Register terminal cleanup handlers
self._register_cleanup()
def status(
self,
status: RenderableType,
@@ -283,6 +291,38 @@ class ComfyConsole(Console):
return status_renderable
def _register_cleanup(self):
"""Register terminal cleanup handlers."""
if not ComfyConsole._cleanup_registered:
ComfyConsole._cleanup_registered = True
# Register cleanup on normal exit
atexit.register(self._cleanup_terminal)
# Register cleanup on signals
signal.signal(signal.SIGINT, self._signal_handler)
signal.signal(signal.SIGTERM, self._signal_handler)
def _cleanup_terminal(self):
"""Restore terminal to a clean state."""
try:
# Show cursor using ANSI escape codes
sys.stdout.write("\x1b[?25h") # Show cursor
sys.stdout.write("\x1b[0m") # Reset attributes
sys.stdout.flush()
# Also use Rich's method
self.show_cursor(True)
except Exception:
# Silently fail if cleanup fails
pass
def _signal_handler(self, signum, frame):
"""Handle signals with cleanup."""
self._cleanup_terminal()
# Exit after cleanup
sys.exit(1)
catppuccin_mocha = {
# Colors based on "CatppuccinMocha" from Gogh themes