Print colored exception tracebacks in interactive terminals

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)

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,
         )