import sys
import os
from pathlib import Path
from functools import wraps
from loguru import logger
from rich.progress import Progress, SpinnerColumn, TextColumn
def setup_logger(log_file):
"""Set up the logger with a file and console handler."""
os.makedirs(Path(log_file).parent, exist_ok=True)
logger.remove()
logger.add(
sys.stderr,
format="{time:YYYY-MM-DD HH:mm:ss} | {level: <8} | {name}:{function}:{line} - {message}",
)
logger.add(log_file, rotation="1MB", retention="7 days", encoding="utf-8")
def task_wrapper(func):
"""Wrapper to log the start and end of a task."""
@wraps(func)
def wrapper(*args, **kwargs):
func_name = func.__name__
logger.info(f"Starting {func_name}")
try:
result = func(*args, **kwargs)
logger.info(f"Finished {func_name}")
return result
except Exception as e:
logger.exception(f"Error in {func_name}: {str(e)}")
raise
return wrapper
def get_rich_progress():
"""Get a Rich Progress object."""
return Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
transient=True,
)