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() DOWNLOAD_LICENCE_ONLY.set()
try: 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: with ThreadPoolExecutor(downloads) as pool:
for download in futures.as_completed( for download in futures.as_completed(
( (
@@ -1009,7 +1010,7 @@ class dl:
multiplex_tasks.append((task_id, task_tracks)) 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: for task_id, task_tracks in multiplex_tasks:
progress.start_task(task_id) # TODO: Needed? progress.start_task(task_id) # TODO: Needed?
muxed_path, return_code, errors = task_tracks.mux( muxed_path, return_code, errors = task_tracks.mux(

View File

@@ -1,4 +1,7 @@
import atexit
import logging import logging
import signal
import sys
from datetime import datetime from datetime import datetime
from types import ModuleType from types import ModuleType
from typing import IO, Callable, Iterable, List, Literal, Mapping, Optional, Union from typing import IO, Callable, Iterable, List, Literal, Mapping, Optional, Union
@@ -167,6 +170,8 @@ class ComfyConsole(Console):
time.monotonic. time.monotonic.
""" """
_cleanup_registered = False
def __init__( def __init__(
self, self,
*, *,
@@ -233,6 +238,9 @@ class ComfyConsole(Console):
if log_renderer: if log_renderer:
self._log_render = log_renderer self._log_render = log_renderer
# Register terminal cleanup handlers
self._register_cleanup()
def status( def status(
self, self,
status: RenderableType, status: RenderableType,
@@ -283,6 +291,38 @@ class ComfyConsole(Console):
return status_renderable 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 = { catppuccin_mocha = {
# Colors based on "CatppuccinMocha" from Gogh themes # Colors based on "CatppuccinMocha" from Gogh themes