Add log formatter with colour (#277)

* Add log formatter with colour

pyproject.toml

@@ -64,6 +64,7 @@ dev = [
     "tox-direct",
     "types-mock",
     "types-pyyaml",
+    "colorlog"
 ]
 
 [project.scripts]

src/ophyd_async/log.py

@@ -0,0 +1,130 @@
+import logging
+import sys
+
+import colorlog
+
+__all__ = (
+    "config_ophyd_async_logging",
+    "logger",
+    "set_handler",
+)
+
+DEFAULT_FORMAT = (
+    "%(log_color)s[%(levelname)1.1s %(asctime)s.%(msecs)03d "
+    "%(module)s:%(lineno)d] %(message)s"
+)
+
+DEFAULT_DATE_FORMAT = "%y%m%d %H:%M:%S"
+
+DEFAULT_LOG_COLORS = {
+    "DEBUG": "cyan",
+    "INFO": "green",
+    "WARNING": "yellow",
+    "ERROR": "red",
+    "CRITICAL": "red,bg_white",
+}
+
+
+class ColoredFormatterWithDeviceName(colorlog.ColoredFormatter):
+    def format(self, record):
+        message = super().format(record)
+        if hasattr(record, "ophyd_async_device_name"):
+            message = f"[{record.ophyd_async_device_name}]{message}"
+        return message
+
+
+def _validate_level(level) -> int:
+    """
+    Return an int for level comparison
+    """
+    if isinstance(level, int):
+        levelno = level
+    elif isinstance(level, str):
+        levelno = logging.getLevelName(level)
+
+    if isinstance(levelno, int):
+        return levelno
+    else:
+        raise ValueError(
+            "Your level is illegal, please use "
+            "'CRITICAL', 'FATAL', 'ERROR', 'WARNING', 'INFO', or 'DEBUG'."
+        )
+
+
+logger = logging.getLogger("ophyd_async")
+
+current_handler = None  # overwritten below
+
+
+def config_ophyd_async_logging(
+    file=sys.stdout,
+    fmt=DEFAULT_FORMAT,
+    datefmt=DEFAULT_DATE_FORMAT,
+    color=True,
+    level="WARNING",
+):
+    """
+    Set a new handler on the ``logging.getLogger('ophyd_async')`` logger.
+    If this is called more than once, the handler from the previous invocation
+    is removed (if still present) and replaced.
+
+    Parameters
+    ----------
+    file : object with ``write`` method or filename string
+        Default is ``sys.stdout``.
+    fmt : Overall logging format
+    datefmt : string
+        Date format. Default is ``'%H:%M:%S'``.
+    color : boolean
+        Use ANSI color codes. True by default.
+    level : str or int
+        Python logging level, given as string or corresponding integer.
+        Default is 'WARNING'.
+    Returns
+    -------
+    handler : logging.Handler
+        The handler, which has already been added to the 'ophyd_async' logger.
+    Examples
+    --------
+    Log to a file.
+        config_ophyd_async_logging(file='/tmp/what_is_happening.txt')
+    Include the date along with the time. (The log messages will always include
+    microseconds, which are configured separately, not as part of 'datefmt'.)
+        config_ophyd_async_logging(datefmt="%Y-%m-%d %H:%M:%S")
+    Turn off ANSI color codes.
+        config_ophyd_async_logging(color=False)
+    Increase verbosity: show level DEBUG or higher.
+        config_ophyd_async_logging(level='DEBUG')
+    """
+    global current_handler
+
+    if isinstance(file, str):
+        handler = logging.FileHandler(file)
+        formatter = ColoredFormatterWithDeviceName(
+            fmt=fmt, datefmt=datefmt, no_color=True
+        )
+    else:
+        handler = colorlog.StreamHandler(file)
+        formatter = ColoredFormatterWithDeviceName(
+            fmt=fmt, datefmt=datefmt, log_colors=DEFAULT_LOG_COLORS, no_color=color
+        )
+
+    levelno = _validate_level(level)
+    handler.setFormatter(formatter)
+    handler.setLevel(levelno)
+
+    if current_handler in logger.handlers:
+        logger.removeHandler(current_handler)
+    logger.addHandler(handler)
+
+    current_handler = handler
+
+    if logger.getEffectiveLevel() > levelno:
+        logger.setLevel(levelno)
+    try:
+        return handler
+    finally:
+        handler.close()
+
+
+set_handler = config_ophyd_async_logging  # for back-compat

tests/test_log.py

@@ -0,0 +1,88 @@
+import io
+import logging
+import logging.handlers
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from ophyd_async import log
+from ophyd_async.core import Device
+from ophyd_async.log import DEFAULT_DATE_FORMAT, DEFAULT_FORMAT
+
+
+def test_validate_level():
+    assert log._validate_level("CRITICAL") == 50
+    assert log._validate_level("ERROR") == 40
+    assert log._validate_level("WARNING") == 30
+    assert log._validate_level("INFO") == 20
+    assert log._validate_level("DEBUG") == 10
+    assert log._validate_level("NOTSET") == 0
+    assert log._validate_level(123) == 123
+    with pytest.raises(ValueError):
+        log._validate_level("MYSTERY")
+
+
+@patch("ophyd_async.log.current_handler")
+@patch("ophyd_async.log.logging.Logger.addHandler")
+def test_default_config_ophyd_async_logging(mock_add_handler, mock_current_handler):
+    log.config_ophyd_async_logging()
+    assert isinstance(log.current_handler, logging.StreamHandler)
+    assert log.logger.getEffectiveLevel() <= logging.WARNING
+
+
+@patch("ophyd_async.log.current_handler")
+@patch("ophyd_async.log.logging.FileHandler")
+@patch("ophyd_async.log.logging.Logger.addHandler")
+def test_config_ophyd_async_logging_with_file_handler(
+    mock_add_handler, mock_file_handler, mock_current_handler
+):
+    log.config_ophyd_async_logging(file="file")
+    assert isinstance(log.current_handler, MagicMock)
+    assert log.logger.getEffectiveLevel() <= logging.WARNING
+
+
+@patch("ophyd_async.log.current_handler")
+def test_config_ophyd_async_logging_removes_extra_handlers(mock_current_handler):
+    # Protect global variable in other pytests
+    class FakeLogger:
+        def __init__(self):
+            self.handlers = []
+            self.removeHandler = MagicMock()
+            self.setLevel = MagicMock()
+
+        def addHandler(self, handler):
+            self.handlers.append(handler)
+
+        def getEffectiveLevel(self):
+            return 100000
+
+    fake_logger = FakeLogger()
+    with (
+        patch("ophyd_async.log.logger", fake_logger),
+    ):
+        log.config_ophyd_async_logging()
+        fake_logger.removeHandler.assert_not_called()
+        log.config_ophyd_async_logging()
+        fake_logger.removeHandler.assert_called()
+
+
+# Full format looks like:
+#'[test_device][W 240501 13:28:08.937 test_log:35] here is a warning\n'
+def test_logger_adapter_ophyd_async_device():
+    log_buffer = io.StringIO()
+    log_stream = logging.StreamHandler(stream=log_buffer)
+    log_stream.setFormatter(
+        log.ColoredFormatterWithDeviceName(
+            fmt=DEFAULT_FORMAT, datefmt=DEFAULT_DATE_FORMAT, no_color=True
+        )
+    )
+    log.logger.addHandler(log_stream)
+
+    device = Device(name="test_device")
+    device.log = logging.LoggerAdapter(
+        logging.getLogger("ophyd_async.devices"),
+        {"ophyd_async_device_name": device.name},
+    )
+    device.log.warning("here is a warning")
+    assert log_buffer.getvalue().startswith("[test_device]")
+    assert log_buffer.getvalue().endswith("here is a warning\n")