Add the Kedro RichHandler

docs/source/logging/logging.md

@@ -19,7 +19,7 @@ We now give some common examples of how you might like to change your project's
 
 ### Using `KEDRO_LOGGING_CONFIG` environment variable
 
-`KEDRO_LOGGING_CONFIG` is an optional environment variable that you can use to specify the path of your logging configuration file, overriding the default path of `conf/base/logging.yml`.
+`KEDRO_LOGGING_CONFIG` is an optional environment variable that you can use to specify the path of your logging configuration file, overriding the default Kedro's `default_logging.yml`.
 
 To use this environment variable, set it to the path of your desired logging configuration file before running any Kedro commands. For example, if you have a logging configuration file located at `/path/to/logging.yml`, you can set `KEDRO_LOGGING_CONFIG` as follows:
 
@@ -43,6 +43,25 @@ Alternatively, if you would like to keep other configuration in `conf/base/loggi
 +  handlers: [console]
 ```
 
+### Customize the `rich` Handler
+
+Kedro's `kedro.extras.logging.RichHandler` is a subclass of [`rich.logging.RichHand[label](vscode-file://vscode-a section of the class.
+
+pp/Applications/Visual%20Studio%20Code.app/Contents/Resources/app/out/vs/code/electron-sandbox/workbench/workbench.html)ler`](https://rich.readthedocs.io/en/stable/reference/logging.html#rich.logging.RichHandler) and supports the same set of arguments. By default, `rich_tracebacks` is set to `True` to use `rich` to render exceptions. However, you can disable it by setting `rich_tracebacks: False`.
+
+When `rich_tracebacks` is set to `True`, the configuration is propagated to [`rich.traceback.install`](https://rich.readthedocs.io/en/stable/reference/traceback.html#rich.traceback.install). If an argument is compatible with `rich.traceback.install`, it will be passed to the traceback's settings.
+
+For instance, you can enable the display of local variables inside `logging.yml` to aid with debugging.
+
+```yaml
+rich:
+  class: kedro.extras.logging.RichHandler
+  rich_tracebacks: True
+  tracebacks_show_locals: True
+```
+
+A comprehensive list of available options can be found in the [RichHandler documentation](https://rich.readthedocs.io/en/stable/reference/logging.html#rich.logging.RichHandler).
+
 ### Use plain console logging
 
 To use plain rather than rich logging, swap the `rich` handler for the `console` one as follows:

features/steps/test_starter/{{ cookiecutter.repo_name }}/conf/base/logging.yml

@@ -34,7 +34,9 @@ handlers:
     delay: True
 
   rich:
-    class: rich.logging.RichHandler
+    class: kedro.logging.RichHandler
+    rich_tracebacks: True
+    # tracebacks_show_locals: False
 
 loggers:
   kedro:

kedro/__init__.py

@@ -4,8 +4,3 @@
 """
 
 __version__ = "0.18.8"
-
-
-import logging
-
-logging.getLogger(__name__).addHandler(logging.NullHandler())

kedro/framework/project/__init__.py

@@ -7,7 +7,6 @@
 import logging.config
 import operator
 import os
-import sys
 import traceback
 import types
 import warnings
@@ -16,10 +15,7 @@
 from pathlib import Path
 from typing import Any
 
-import click
 import importlib_resources
-import rich.pretty
-import rich.traceback
 import yaml
 from dynaconf import LazySettings
 from dynaconf.validator import ValidationError, Validator
@@ -221,18 +217,6 @@ def __init__(self):
         )
         logging_config = Path(path).read_text(encoding="utf-8")
         self.configure(yaml.safe_load(logging_config))
-        logging.captureWarnings(True)
-
-        # We suppress click here to hide tracebacks related to it conversely,
-        # kedro is not suppressed to show its tracebacks for easier debugging.
-        # sys.executable is used to get the kedro executable path to hide the
-        # top level traceback.
-        # Rich traceback handling does not work on databricks. Hopefully this will be
-        # fixed on their side at some point, but until then we disable it.
-        # See https://github.com/Textualize/rich/issues/2455
-        if "DATABRICKS_RUNTIME_VERSION" not in os.environ:
-            rich.traceback.install(suppress=[click, str(Path(sys.executable).parent)])
-        rich.pretty.install()
 
     def configure(self, logging_config: dict[str, Any]) -> None:
         """Configure project logging using ``logging_config`` (e.g. from project

kedro/framework/project/default_logging.yml

@@ -4,7 +4,9 @@ disable_existing_loggers: False
 
 handlers:
   rich:
-    class: rich.logging.RichHandler
+    class: kedro.logging.RichHandler
+    rich_tracebacks: True
+    # tracebacks_show_locals: False
 
 loggers:
   kedro:

kedro/logging.py

@@ -0,0 +1,60 @@
+"""
+This module contains a logging handler class which produces coloured logs and tracebacks.
+"""
+
+import logging
+import os
+import sys
+from pathlib import Path
+
+import click
+import rich.logging
+import rich.pretty
+import rich.traceback
+
+logger = logging.getLogger(__file__)
+
+
+class RichHandler(rich.logging.RichHandler):
+    """Identical to rich's logging handler but with a few extra behaviours:
+    * warnings issued by the `warnings` module are redirected to logging
+    * pretty printing is enabled on the Python REPL (including IPython and Jupyter)
+    * all tracebacks are handled by rich when rich_tracebacks=True
+    * constructor's arguments are mapped and pass to `rich.traceback.install`
+
+    The list of available options of ``RichHandler`` can be found here:
+    https://rich.readthedocs.io/en/stable/reference/logging.html#rich.logging.RichHandler
+
+    The list of available options of `rich.traceback.install` can be found here:
+    https://rich.readthedocs.io/en/stable/reference/traceback.html#rich.traceback.install
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        logging.captureWarnings(True)
+        rich.pretty.install()
+
+        # We suppress click here to hide tracebacks related to it conversely,
+        # kedro is not suppressed to show its tracebacks for easier debugging.
+        # sys.executable is used to get the kedro executable path to hide the
+        # top level traceback.
+
+        traceback_install_kwargs = {
+            "suppress": [click, str(Path(sys.executable).parent)]
+        }
+
+        # Mapping Arguments from RichHandler's Constructor to rich.traceback.install
+        prefix = "tracebacks_"
+        for key, value in kwargs.items():
+            if key.startswith(prefix):
+                key_prefix_removed = key[len(prefix) :]
+                if key_prefix_removed == "suppress":
+                    traceback_install_kwargs[key_prefix_removed].extend(value)
+                else:
+                    traceback_install_kwargs[key_prefix_removed] = value
+
+        if self.rich_tracebacks and "DATABRICKS_RUNTIME_VERSION" not in os.environ:
+            # Rich traceback handling does not work on databricks. Hopefully this will be
+            # fixed on their side at some point, but until then we disable it.
+            # See https://github.com/Textualize/rich/issues/2455
+            rich.traceback.install(**traceback_install_kwargs)

kedro/templates/project/{{ cookiecutter.repo_name }}/conf/base/logging.yml

@@ -24,7 +24,9 @@ handlers:
     delay: True
 
   rich:
-    class: rich.logging.RichHandler
+    class: kedro.logging.RichHandler
+    rich_tracebacks: True
+    # tracebacks_show_locals: True
 
 loggers:
   kedro:

tests/framework/project/test_logging.py

@@ -1,4 +1,4 @@
-# pylint: disable=import-outside-toplevel,unused-import,reimported
+# pylint: disable=import-outside-toplevel
 import logging
 import sys
 from pathlib import Path
@@ -11,7 +11,9 @@
 default_logging_config = {
     "version": 1,
     "disable_existing_loggers": False,
-    "handlers": {"rich": {"class": "rich.logging.RichHandler"}},
+    "handlers": {
+        "rich": {"class": "kedro.logging.RichHandler", "rich_tracebacks": True}
+    },
     "loggers": {"kedro": {"level": "INFO"}},
     "root": {"handlers": ["rich"]},
 }
@@ -35,8 +37,9 @@ def test_environment_variable_logging_config(monkeypatch, tmp_path):
     logging_config = {"version": 1, "loggers": {"kedro": {"level": "WARNING"}}}
     with config_path.open("w", encoding="utf-8") as f:
         yaml.dump(logging_config, f)
-    del sys.modules["kedro.framework.project"]
-    from kedro.framework.project import LOGGING  # noqa
+    from kedro.framework.project import _ProjectLogging
+
+    LOGGING = _ProjectLogging()
 
     assert LOGGING.data == logging_config
     assert logging.getLogger("kedro").level == logging.WARNING
@@ -50,24 +53,65 @@ def test_configure_logging():
 
 
 def test_rich_traceback_enabled(mocker):
-    """Note we need to force reload; just doing from kedro.framework.project import ...
-    will not call rich.traceback.install again. Using importlib.reload does not work
-    well with other tests here, hence the manual del sys.modules."""
     rich_traceback_install = mocker.patch("rich.traceback.install")
     rich_pretty_install = mocker.patch("rich.pretty.install")
-    del sys.modules["kedro.framework.project"]
-    from kedro.framework.project import LOGGING  # noqa
+
+    LOGGING.configure(default_logging_config)
 
     rich_traceback_install.assert_called()
     rich_pretty_install.assert_called()
 
 
+def test_rich_traceback_not_installed(mocker):
+    rich_traceback_install = mocker.patch("rich.traceback.install")
+    rich_pretty_install = mocker.patch("rich.pretty.install")
+    rich_handler = {
+        "class": "kedro.logging.RichHandler",
+        "rich_tracebacks": False,
+        "tracebacks_show_locals": True,
+    }
+    test_logging_config = default_logging_config.copy()
+    test_logging_config["handlers"]["rich"] = rich_handler
+
+    LOGGING.configure(test_logging_config)
+
+    rich_pretty_install.assert_called_once()
+    rich_traceback_install.assert_not_called()
+
+
+def test_rich_traceback_configuration(mocker):
+    rich_traceback_install = mocker.patch("rich.traceback.install")
+    rich_pretty_install = mocker.patch("rich.pretty.install")
+
+    import click
+
+    sys_executable_path = str(Path(sys.executable).parent)
+    traceback_install_defaults = {"suppress": [click, sys_executable_path]}
+    fake_path = "dummy"
+    rich_handler = {
+        "class": "kedro.logging.RichHandler",
+        "rich_tracebacks": True,
+        "tracebacks_show_locals": True,
+        "tracebacks_suppress": [fake_path],
+    }
+
+    test_logging_config = default_logging_config.copy()
+    test_logging_config["handlers"]["rich"] = rich_handler
+    LOGGING.configure(test_logging_config)
+
+    expected_install_defaults = traceback_install_defaults.copy()
+    expected_install_defaults["suppress"].extend([fake_path])
+    expected_install_defaults["show_locals"] = True
+    rich_traceback_install.assert_called_with(**expected_install_defaults)
+    rich_pretty_install.assert_called_once()
+
+
 def test_rich_traceback_disabled_on_databricks(mocker, monkeypatch):
     monkeypatch.setenv("DATABRICKS_RUNTIME_VERSION", "1")
     rich_traceback_install = mocker.patch("rich.traceback.install")
     rich_pretty_install = mocker.patch("rich.pretty.install")
-    del sys.modules["kedro.framework.project"]
-    from kedro.framework.project import LOGGING  # noqa
+
+    LOGGING.configure(default_logging_config)
 
     rich_traceback_install.assert_not_called()
     rich_pretty_install.assert_called()