Add the Kedro RichHandler (#2592)

* Add the Kedro RichHandler Signed-off-by: Nok Chan <[email protected]> * Clean up imports and comments Signed-off-by: Nok Chan <[email protected]> * Map arguments to add flexibility to customize arguments available in RichHandler Signed-off-by: Nok Chan <[email protected]> * Clean up rich_logger.py and unused imports Signed-off-by: Nok Chan <[email protected]> * Add docs - first iteration Signed-off-by: Nok Chan <[email protected]> * Fix imports and small refactoring Signed-off-by: Nok Chan <[email protected]> * Improve the doc Signed-off-by: Nok Chan <[email protected]> * Update the project's template Signed-off-by: Nok Chan <[email protected]> * format docstring Signed-off-by: Nok Chan <[email protected]> * Apply Jo's comment on the doc Signed-off-by: Nok Chan <[email protected]> * Fix the link with RichHandler Signed-off-by: Nok Chan <[email protected]> * Apply suggestions from code review Co-authored-by: Antony Milne <[email protected]> * Fix broken doc Signed-off-by: Nok Chan <[email protected]> * update template defaults Signed-off-by: Nok Chan <[email protected]> * refactoring RichHandler and improve doc Signed-off-by: Nok Chan <[email protected]> * Fix broken test due to new default RichHandler Signed-off-by: Nok Chan <[email protected]> * Fix doc Signed-off-by: Nok Chan <[email protected]> * Move kedro.extras.logging.RichHandler -> kedro.logging.RichHandler Signed-off-by: Nok Chan <[email protected]> * add kedro/logging.py and remove the unnecessary import that cause namespace collision Signed-off-by: Nok Chan <[email protected]> * Refactor test Signed-off-by: Nok Chan <[email protected]> * Add more tests Signed-off-by: Nok Chan <[email protected]> * Fix lots of import, docstring Signed-off-by: Nok Chan <[email protected]> * Reorder docstring Signed-off-by: Nok Chan <[email protected]> * Fix tests Signed-off-by: Nok Chan <[email protected]> * Link options to doc Signed-off-by: Nok Chan <[email protected]> * Clean up doc and unused logger Signed-off-by: Nok Chan <[email protected]> * Fix spelling `customise` Signed-off-by: Nok Chan <[email protected]> * Refactor test with fixture Signed-off-by: Nok Chan <[email protected]> * Remove unnecessary copy() because it is now a fixture Signed-off-by: Nok Chan <[email protected]> * Split up test_rich_traceback_configuration() to smaller tests Signed-off-by: Nok Chan <[email protected]> * Add RELEASE.md Signed-off-by: Nok Chan <[email protected]> * Update doc Signed-off-by: Nok Chan <[email protected]> * Fix template logging
kedro-org · May 30, 2023 · 4a4bd73 · 4a4bd73
1 parent 2f248a4
commit 4a4bd73
Show file tree

Hide file tree

Showing 12 changed files with 225 additions and 54 deletions.
diff --git a/RELEASE.md b/RELEASE.md
@@ -13,6 +13,7 @@
 ## Major features and improvements
 * `kedro run --params` now updates interpolated parameters correctly when using `OmegaConfigLoader`.
 * Added `metadata` attribute to `kedro.io` datasets. This is ignored by Kedro, but may be consumed by users or external plugins.
+* Added `kedro.logging.RichHandler`. This replaces the default `rich.logging.RichHandler` and is more flexible, user can turn off the `rich` traceback if needed.
 
 ## Bug fixes and other changes
 * `OmegaConfigLoader` will return a `dict` instead of `DictConfig`.
@@ -141,12 +142,12 @@ Many thanks to the following Kedroids for contributing PRs to this release:
 * You can configure config file patterns through `settings.py` without creating a custom config loader.
 * Added the following new datasets:
 
-| Type                                 | Description                                                                | Location                      |
-| ------------------------------------ | -------------------------------------------------------------------------- | ----------------------------- |
-| `svmlight.SVMLightDataSet` | Work with svmlight/libsvm files using scikit-learn library | `kedro.extras.datasets.svmlight` |
-| `video.VideoDataSet`                 | Read and write video files from a filesystem                               | `kedro.extras.datasets.video` |
-| `video.video_dataset.SequenceVideo`  | Create a video object from an iterable sequence to use with `VideoDataSet` | `kedro.extras.datasets.video` |
-| `video.video_dataset.GeneratorVideo` | Create a video object from a generator to use with `VideoDataSet`          | `kedro.extras.datasets.video` |
+| Type                                 | Description                                                                | Location                         |
+| ------------------------------------ | -------------------------------------------------------------------------- | -------------------------------- |
+| `svmlight.SVMLightDataSet`           | Work with svmlight/libsvm files using scikit-learn library                 | `kedro.extras.datasets.svmlight` |
+| `video.VideoDataSet`                 | Read and write video files from a filesystem                               | `kedro.extras.datasets.video`    |
+| `video.video_dataset.SequenceVideo`  | Create a video object from an iterable sequence to use with `VideoDataSet` | `kedro.extras.datasets.video`    |
+| `video.video_dataset.GeneratorVideo` | Create a video object from a generator to use with `VideoDataSet`          | `kedro.extras.datasets.video`    |
 * Implemented support for a functional definition of schema in `dask.ParquetDataSet` to work with the `dask.to_parquet` API.
 
 ## Bug fixes and other changes

diff --git a/docs/source/conf.py b/docs/source/conf.py
@@ -191,7 +191,7 @@
 here = Path(__file__).parent.absolute()
 html_logo = str(here / "kedro_logo.svg")
 
-# Theme options are theme-specific and customize the look and feel of a theme
+# Theme options are theme-specific and customise the look and feel of a theme
 # further.  For a list of options available for each theme, see the
 # documentation.
 #

diff --git a/docs/source/kedro.logging.RichHandler.rst b/docs/source/kedro.logging.RichHandler.rst
@@ -0,0 +1,6 @@
+kedro.logging.RichHandler
+=========================
+
+.. currentmodule:: kedro.logging
+
+.. autoclass:: RichHandler
diff --git a/docs/source/kedro.logging.rst b/docs/source/kedro.logging.rst
@@ -0,0 +1,20 @@
+kedro.logging
+=============
+
+.. rubric:: Description
+
+.. automodule:: kedro.logging
+
+
+
+
+
+
+
+   .. rubric:: Classes
+
+   .. autosummary::
+      :toctree:
+      :template: autosummary/class.rst
+
+      RichHandler
diff --git a/docs/source/logging/logging.md b/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:
 
@@ -30,7 +30,7 @@ export KEDRO_LOGGING_CONFIG=/path/to/logging.yml
 After setting the environment variable, any subsequent Kedro commands will use the logging configuration file at the specified path.
 
 ```{note}
-If the `KEDRO_LOGGING_CONFIG` environment variable is not set, Kedro will default to using the logging configuration file at the project's default location of `conf/base/logging.yml`.
+If the `KEDRO_LOGGING_CONFIG` environment variable is not set, Kedro will default to using the logging configuration file at the project's default location of  Kedro's `default_logging.yml`.
 ```
 ### Disable file-based logging
 
@@ -43,6 +43,28 @@ Alternatively, if you would like to keep other configuration in `conf/base/loggi
 +  handlers: [console]
 ```
 
+### Customise the `rich` Handler
+
+Kedro's `kedro.extras.logging.RichHandler` is a subclass of [`rich.logging.RichHandler`](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`.
+
+```{note}
+If you want to disable `rich`'s tracebacks, you must set `KEDRO_LOGGING_CONFIG` to point to your local config i.e. `conf/base/logging.yml`.
+```
+
+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:

diff --git a/features/steps/test_starter/{{ cookiecutter.repo_name }}/conf/base/logging.yml b/features/steps/test_starter/{{ cookiecutter.repo_name }}/conf/base/logging.yml
@@ -34,7 +34,11 @@ handlers:
     delay: True
 
   rich:
-    class: rich.logging.RichHandler
+    class: kedro.logging.RichHandler
+    rich_tracebacks: True
+    # Advance options for customisation.
+    # See https://docs.kedro.org/en/stable/logging/logging.html#project-side-logging-configuration
+    # tracebacks_show_locals: False
 
 loggers:
   kedro:

diff --git a/kedro/__init__.py b/kedro/__init__.py
@@ -4,8 +4,3 @@
 """
 
 __version__ = "0.18.8"
-
-
-import logging
-
-logging.getLogger(__name__).addHandler(logging.NullHandler())
diff --git a/kedro/framework/project/__init__.py b/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

diff --git a/kedro/framework/project/default_logging.yml b/kedro/framework/project/default_logging.yml
@@ -4,7 +4,11 @@ disable_existing_loggers: False
 
 handlers:
   rich:
-    class: rich.logging.RichHandler
+    class: kedro.logging.RichHandler
+    rich_tracebacks: True
+    # Advance options for customisation.
+    # See https://docs.kedro.org/en/stable/logging/logging.html#project-side-logging-configuration
+    # tracebacks_show_locals: False
 
 loggers:
   kedro:

diff --git a/kedro/logging.py b/kedro/logging.py
@@ -0,0 +1,58 @@
+"""
+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
+
+
+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 passed 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)
diff --git a/kedro/templates/project/{{ cookiecutter.repo_name }}/conf/base/logging.yml b/kedro/templates/project/{{ cookiecutter.repo_name }}/conf/base/logging.yml
@@ -24,7 +24,11 @@ handlers:
     delay: True
 
   rich:
-    class: rich.logging.RichHandler
+    class: kedro.logging.RichHandler
+    rich_tracebacks: True
+    # Advance options for customisation.
+    # See https://docs.kedro.org/en/stable/logging/logging.html#project-side-logging-configuration
+    # tracebacks_show_locals: False
 
 loggers:
   kedro: