112 lines
3.5 KiB
Python
112 lines
3.5 KiB
Python
"""
|
|
DICOM file cleanup utility to ensure temporary DICOM files are removed even on script termination.
|
|
"""
|
|
import os
|
|
import sys
|
|
import shutil
|
|
import atexit
|
|
import signal
|
|
from config import settings
|
|
from utils.logger import main_logger as logger
|
|
|
|
# Global registry of directories to clean up
|
|
_cleanup_dirs = set()
|
|
|
|
def register_cleanup_dir(directory):
|
|
"""
|
|
Register a directory for cleanup when the script exits.
|
|
|
|
Args:
|
|
directory (str): Path to the directory to be cleaned up
|
|
"""
|
|
global _cleanup_dirs
|
|
if os.path.exists(directory) and os.path.isdir(directory):
|
|
_cleanup_dirs.add(directory)
|
|
logger.debug(f"Registered directory for cleanup: {directory}")
|
|
|
|
def cleanup_dicom_files():
|
|
"""
|
|
Remove all registered DICOM directories.
|
|
This function is called when the script exits, ensuring cleanup even on abnormal termination.
|
|
"""
|
|
global _cleanup_dirs
|
|
logger.info(f"Running cleanup for {len(_cleanup_dirs)} directories")
|
|
|
|
for directory in _cleanup_dirs:
|
|
try:
|
|
if os.path.exists(directory):
|
|
shutil.rmtree(directory)
|
|
logger.info(f"Cleanup: Removed DICOM files from {directory}")
|
|
except Exception as e:
|
|
logger.error(f"Cleanup: Failed to remove DICOM files from {directory}: {str(e)}")
|
|
|
|
# Clear the registry after cleanup
|
|
_cleanup_dirs.clear()
|
|
|
|
def cleanup_study_directory(study_uid):
|
|
"""
|
|
Immediately remove the directory for a specific study.
|
|
|
|
Args:
|
|
study_uid (str): Study Instance UID of the study to clean up
|
|
|
|
Returns:
|
|
bool: True if cleanup was successful, False otherwise
|
|
"""
|
|
study_dir = os.path.join(settings.DICOM_STORE_DIR, study_uid)
|
|
|
|
if not os.path.exists(study_dir):
|
|
logger.debug(f"Cleanup: Study directory does not exist: {study_dir}")
|
|
return True
|
|
|
|
try:
|
|
shutil.rmtree(study_dir)
|
|
logger.info(f"Cleanup: Removed DICOM files from {study_dir}")
|
|
|
|
# Also remove from the registry if it's there
|
|
global _cleanup_dirs
|
|
if study_dir in _cleanup_dirs:
|
|
_cleanup_dirs.remove(study_dir)
|
|
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Cleanup: Failed to remove DICOM files from {study_dir}: {str(e)}")
|
|
return False
|
|
|
|
def register_exit_handlers():
|
|
"""
|
|
Register cleanup handlers for various exit scenarios.
|
|
"""
|
|
# Register for normal exit
|
|
atexit.register(cleanup_dicom_files)
|
|
|
|
# Register for signals
|
|
signal.signal(signal.SIGINT, signal_handler) # Ctrl+C
|
|
signal.signal(signal.SIGTERM, signal_handler) # Termination request
|
|
|
|
logger.info("Registered cleanup handlers for script exit")
|
|
|
|
def signal_handler(sig, frame):
|
|
"""
|
|
Handle termination signals by performing cleanup before exit.
|
|
|
|
Args:
|
|
sig: Signal number
|
|
frame: Current stack frame
|
|
"""
|
|
signal_name = {
|
|
signal.SIGINT: 'SIGINT (Ctrl+C)',
|
|
signal.SIGTERM: 'SIGTERM'
|
|
}.get(sig, f'Signal {sig}')
|
|
|
|
logger.info(f"Received {signal_name}, cleaning up...")
|
|
cleanup_dicom_files()
|
|
# Don't call exit() directly in a web application context
|
|
# Let the calling application handle the exit process
|
|
# This fix makes it compatible with FastAPI and other web frameworks
|
|
# Raise SystemExit instead of calling exit() directly
|
|
if hasattr(frame, '_is_web_app') and frame._is_web_app:
|
|
return
|
|
else:
|
|
# For command-line applications, we can exit directly
|
|
sys.exit(0) |