Basic logging utilities and filtering

supervision/utils/logging.py

@@ -0,0 +1,49 @@
+import logging
+from datetime import datetime, timedelta
+from typing import Dict, Tuple
+
+
+def setup_logger(**levels: int):
+    formatter = logging.Formatter(
+        "%(asctime)s %(name)s:%(lineno)d %(levelname)s %(message)s"
+    )
+
+    console_handler = logging.StreamHandler()
+    console_handler.setLevel(logging.DEBUG)
+    console_handler.setFormatter(formatter)
+
+    root = logging.getLogger("supervision")
+    root.setLevel(logging.WARNING)
+    root.addHandler(console_handler)
+
+    for logger_name, logger_level in levels.items():
+        logging.getLogger(logger_name).setLevel(logger_level)
+
+
+_LogId = Tuple[str, int]
+
+
+def _get_record_id(record: logging.LogRecord) -> _LogId:
+    return record.pathname, record.lineno
+
+
+class TimeBetweenOccurrencesFilter(logging.Filter):
+    """
+    Adds filtering based on time elapsed between two same logging calls.
+    Useful to prevent displaying too much messages, e.g. in the loop.
+    """
+
+    def __init__(self, min_interval: timedelta):
+        super().__init__()
+        self._min_interval = min_interval
+
+        self._last_activity: Dict[_LogId, datetime] = {}
+
+    def filter(self, record: logging.LogRecord) -> bool:
+        record_id = _get_record_id(record)
+        now = datetime.fromtimestamp(record.created)
+        last_activity = self._last_activity.get(record_id)
+        self._last_activity[record_id] = now
+        if last_activity is not None:
+            return now - last_activity >= self._min_interval
+        return True

test/utils/test_logging.py

@@ -0,0 +1,47 @@
+import logging
+import time
+from datetime import timedelta
+
+import pytest
+
+from supervision.utils.logging import TimeBetweenOccurrencesFilter, setup_logger
+
+
+def test_logging_without_setup_logger(capsys: pytest.CaptureFixture[str]):
+    logger = logging.getLogger("supervision.test_logging_without_setup_logger")
+    logger.warning("Info message")
+    captured = capsys.readouterr()
+    assert captured.out == ""
+    assert captured.err == ""
+
+
+def test_setup_logger(capsys: pytest.CaptureFixture[str]):
+    setup_logger()
+    logger = logging.getLogger("supervision.test_setup_logger")
+    logger.warning("Info message")
+    captured = capsys.readouterr()
+    assert captured.out == ""
+    assert "Info message" in captured.err
+
+
+def function_that_logs(logger: logging.Logger, i: int):
+    logger.warning("Info message %d", i)
+
+
+def test_time_between_occurrences_filter(capsys: pytest.CaptureFixture[str]):
+    setup_logger()
+    logger = logging.getLogger("supervision.test_time_between_occurrences_filter")
+    logger.addFilter(TimeBetweenOccurrencesFilter(timedelta(milliseconds=100)))
+
+    function_that_logs(logger, 0)
+    captured = capsys.readouterr()
+    assert "Info message 0" in captured.err
+
+    function_that_logs(logger, 1)
+    captured = capsys.readouterr()
+    assert "Info message 1" not in captured.err
+
+    time.sleep(0.1)
+    function_that_logs(logger, 2)
+    captured = capsys.readouterr()
+    assert "Info message 2" in captured.err