diff --git a/datapipe/cli.py b/datapipe/cli.py index 860527e6..75206c72 100644 --- a/datapipe/cli.py +++ b/datapipe/cli.py @@ -112,19 +112,35 @@ def cli( def setup_logging(): import logging - if debug: - datapipe_logger = logging.getLogger("datapipe") - datapipe_logger.setLevel(logging.DEBUG) - - datapipe_core_steps_logger = logging.getLogger("datapipe.core_steps") - datapipe_core_steps_logger.setLevel(logging.DEBUG) + from rich.console import Console + from rich.logging import RichHandler + from traceback_with_variables.color import supports_ansi - logging.basicConfig(level=logging.DEBUG, stream=sys.stderr) + if debug: + log_level = logging.DEBUG + logging.getLogger("datapipe").setLevel(logging.DEBUG) + logging.getLogger("datapipe.core_steps").setLevel(logging.DEBUG) else: + log_level = logging.INFO for logger_name in [None, "datapipe", "tqdm_loggable"]: - logger = logging.getLogger(logger_name) - logger.setLevel(logging.INFO) - logging.basicConfig(level=logging.INFO, stream=sys.stderr) + logging.getLogger(logger_name).setLevel(logging.INFO) + + if supports_ansi(sys.stderr): + handler = RichHandler( + console=Console(stderr=True, force_terminal=True), + show_time=False, + show_path=False, + markup=False, + rich_tracebacks=False, + ) + logging.basicConfig(level=log_level, format="%(message)s", handlers=[handler], force=True) + else: + logging.basicConfig( + level=log_level, + stream=sys.stderr, + format="%(levelname)s:%(name)s:%(message)s", + force=True, + ) if debug_sql: logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO) diff --git a/datapipe/event_logger.py b/datapipe/event_logger.py index 35169c2b..9e47cca6 100644 --- a/datapipe/event_logger.py +++ b/datapipe/event_logger.py @@ -1,12 +1,19 @@ import logging +import sys from traceback_with_variables import format_exc +from traceback_with_variables.color import supports_ansi from datapipe.run_config import RunConfig logger = logging.getLogger("datapipe.event_logger") +def _format_exception(exc: Exception) -> str: + # for_file enables ANSI colors when stderr is an interactive terminal. + return format_exc(exc, for_file=sys.stderr) + + class EventLogger: def log_state( self, @@ -31,9 +38,19 @@ def log_error( run_config: RunConfig | None = None, ) -> None: if run_config is not None: - logger.error(f"Error in step {run_config.labels.get('step_name')}: {type} {message}\n{description}") + header = f"Error in step {run_config.labels.get('step_name')}: {type} {message}" + else: + header = f"Error: {type} {message}" + + logger.error(header) + if not description: + return + + # Write colored tracebacks directly to stderr so ANSI escapes are not altered by logging. + if supports_ansi(sys.stderr): + sys.stderr.write(f"{description}\n") else: - logger.error(f"Error: {type} {message}\n{description}") + logger.error(description) def log_exception( self, @@ -43,7 +60,7 @@ def log_exception( self.log_error( type=type(exc).__name__, message=str(exc), - description=format_exc(exc), + description=_format_exception(exc), params=[], # exc.args, # Not all args can be serialized to JSON, dont really need them run_config=run_config, )