Add logging to MESA (#2506)

docs/apis/api_main.md

@@ -10,5 +10,6 @@ space
 datacollection
 batchrunner
 visualization
+logging
 experimental
 ```

docs/apis/mesa_logging.md

@@ -0,0 +1,7 @@
+# logging
+
+```{eval-rst}
+.. automodule:: mesa.mesa_logging
+   :members:
+   :inherited-members:
+```

mesa/examples/basic/boltzmann_wealth_model/app.py

@@ -1,10 +1,13 @@
 from mesa.examples.basic.boltzmann_wealth_model.model import BoltzmannWealth
+from mesa.mesa_logging import DEBUG, log_to_stderr
 from mesa.visualization import (
     SolaraViz,
     make_plot_component,
     make_space_component,
 )
 
+log_to_stderr(DEBUG)
+
 
 def agent_portrayal(agent):
     color = agent.wealth  # we are using a colormap to translate wealth to color

mesa/mesa_logging.py

@@ -0,0 +1,190 @@
+"""This provides logging functionality for MESA.
+
+It is modeled on the default `logging approach that comes with Python <https://docs.python.org/library/logging.html>`_.
+
+
+"""
+
+import inspect
+import logging
+from functools import wraps
+from logging import DEBUG, INFO
+
+__all__ = [
+    "get_rootlogger",
+    "get_module_logger",
+    "log_to_stderr",
+    "DEBUG",
+    "INFO",
+    "DEFAULT_LEVEL",
+    "LOGGER_NAME",
+    "method_logger",
+    "function_logger",
+]
+LOGGER_NAME = "MESA"
+DEFAULT_LEVEL = DEBUG
+
+
+def create_module_logger(name: str | None = None):
+    """Helper function for creating a module logger.
+
+    Args:
+        name (str): The name to be given to the logger. If the name is None, the name defaults to the name of the module.
+
+    """
+    if name is None:
+        frm = inspect.stack()[1]
+        mod = inspect.getmodule(frm[0])
+        name = mod.__name__
+    logger = logging.getLogger(f"{LOGGER_NAME}.{name}")
+
+    _module_loggers[name] = logger
+    return logger
+
+
+def get_module_logger(name: str):
+    """Helper function for getting the module logger.
+
+    Args:
+        name (str): The name of the module in which the method being decorated is located
+
+    """
+    try:
+        logger = _module_loggers[name]
+    except KeyError:
+        logger = create_module_logger(name)
+
+    return logger
+
+
+_rootlogger = None
+_module_loggers = {}
+_logger = get_module_logger(__name__)
+
+
+class MESAColorFormatter(logging.Formatter):
+    """Custom formatter for color based formatting."""
+
+    grey = "\x1b[38;20m"
+    green = "\x1b[32m"
+    yellow = "\x1b[33;20m"
+    red = "\x1b[31;20m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+    format = (
+        "[%(asctime)s] [%(name)s] [%(levelname)s] %(message)s [%(filename)s:%(lineno)d]"
+    )
+
+    FORMATS = {
+        logging.DEBUG: grey + format + reset,
+        logging.INFO: green + format + reset,
+        logging.WARNING: yellow + format + reset,
+        logging.ERROR: red + format + reset,
+        logging.CRITICAL: bold_red + format + reset,
+    }
+
+    def format(self, record):
+        """Format record."""
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
+
+def method_logger(name: str):
+    """Decorator for adding logging to a method.
+
+    Args:
+        name (str): The name of the module in which the method being decorated is located
+
+    """
+    logger = get_module_logger(name)
+    classname = inspect.getouterframes(inspect.currentframe())[1][3]
+
+    def real_decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            # hack, because log is applied to methods, we can get
+            # object instance as first arguments in args
+            logger.debug(
+                f"calling {classname}.{func.__name__} with {args[1::]} and {kwargs}"
+            )
+            res = func(*args, **kwargs)
+            return res
+
+        return wrapper
+
+    return real_decorator
+
+
+def function_logger(name):
+    """Decorator for adding logging to a Function.
+
+    Args:
+        name (str): The name of the module in which the function being decorated is located
+
+    """
+    logger = get_module_logger(name)
+
+    def real_decorator(func):
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            logger.debug(f"calling {func.__name__} with {args} and {kwargs}")
+            res = func(*args, **kwargs)
+            return res
+
+        return wrapper
+
+    return real_decorator
+
+
+def get_rootlogger():
+    """Returns root logger used by MESA.
+
+    Returns:
+        the root logger of MESA
+
+    """
+    global _rootlogger  # noqa: PLW0603
+
+    if not _rootlogger:
+        _rootlogger = logging.getLogger(LOGGER_NAME)
+        _rootlogger.handlers = []
+        _rootlogger.addHandler(logging.NullHandler())
+        _rootlogger.setLevel(DEBUG)
+
+    return _rootlogger
+
+
+def log_to_stderr(level: int | None = None, pass_root_logger_level: bool = False):
+    """Turn on logging and add a handler which prints to stderr.
+
+    Args:
+        level: minimum level of the messages that will be logged
+        pass_root_logger_level: bool, optional. Default False
+                if True, all module loggers will be set to the same logging level as the root logger.
+
+    """
+    if not level:
+        level = DEFAULT_LEVEL
+
+    logger = get_rootlogger()
+
+    # avoid creation of multiple stream handlers for logging to console
+    for entry in logger.handlers:
+        if (isinstance(entry, logging.StreamHandler)) and (
+            isinstance(entry.formatter, MESAColorFormatter)
+        ):
+            return logger
+
+    formatter = MESAColorFormatter()
+    handler = logging.StreamHandler()
+    handler.setLevel(level)
+    handler.setFormatter(formatter)
+    logger.addHandler(handler)
+    logger.propagate = False
+
+    if pass_root_logger_level:
+        for _, mod_logger in _module_loggers.items():
+            mod_logger.setLevel(level)
+
+    return logger

mesa/model.py

@@ -17,11 +17,15 @@
 import numpy as np
 
 from mesa.agent import Agent, AgentSet
+from mesa.mesa_logging import create_module_logger, method_logger
 
 SeedLike = int | np.integer | Sequence[int] | np.random.SeedSequence
 RNGLike = np.random.Generator | np.random.BitGenerator
 
 
+_mesa_logger = create_module_logger()
+
+
 class Model:
     """Base class for models in the Mesa ABM library.
 
@@ -42,6 +46,7 @@ class Model:
 
     """
 
+    @method_logger(__name__)
     def __init__(
         self,
         *args: Any,
@@ -112,6 +117,7 @@ def _wrapped_step(self, *args: Any, **kwargs: Any) -> None:
         """Automatically increments time and steps after calling the user's step method."""
         # Automatically increment time and step counters
         self.steps += 1
+        _mesa_logger.info(f"calling model.step for timestep {self.steps} ")
         # Call the original user-defined step method
         self._user_step(*args, **kwargs)
 
@@ -164,6 +170,9 @@ def register_agent(self, agent):
             )
 
         self._all_agents.add(agent)
+        _mesa_logger.debug(
+            f"registered {agent.__class__.__name__} with agent_id {agent.unique_id}"
+        )
 
     def deregister_agent(self, agent):
         """Deregister the agent with the model.
@@ -178,6 +187,7 @@ def deregister_agent(self, agent):
         del self._agents[agent]
         self._agents_by_type[type(agent)].remove(agent)
         self._all_agents.remove(agent)
+        _mesa_logger.debug(f"deregistered agent with agent_id {agent.unique_id}")
 
     def run_model(self) -> None:
         """Run the model until the end condition is reached.

mesa/visualization/solara_viz.py

@@ -33,14 +33,18 @@
 
 import mesa.visualization.components.altair_components as components_altair
 from mesa.experimental.devs.simulator import Simulator
+from mesa.mesa_logging import create_module_logger, function_logger
 from mesa.visualization.user_param import Slider
 from mesa.visualization.utils import force_update, update_counter
 
 if TYPE_CHECKING:
     from mesa.model import Model
 
+_mesa_logger = create_module_logger()
+
 
 @solara.component
+@function_logger(__name__)
 def SolaraViz(
     model: Model | solara.Reactive[Model],
     components: list[reacton.core.Component]
@@ -200,18 +204,25 @@ async def step():
         step, dependencies=[playing.value, running.value], prefer_threaded=False
     )
 
+    @function_logger(__name__)
     def do_step():
         """Advance the model by one step."""
         model.value.step()
         running.value = model.value.running
         force_update()
 
+    @function_logger(__name__)
     def do_reset():
         """Reset the model to its initial state."""
         playing.value = False
         running.value = True
+        _mesa_logger.log(
+            10,
+            f"creating new {model.value.__class__} instance with {model_parameters.value}",
+        )
         model.value = model.value = model.value.__class__(**model_parameters.value)
 
+    @function_logger(__name__)
     def do_play_pause():
         """Toggle play/pause."""
         playing.value = not playing.value
@@ -394,6 +405,7 @@ def ModelCreator(
         [],
     )
 
+    @function_logger(__name__)
     def on_change(name, value):
         model_parameters.value = {**model_parameters.value, name: value}
 

tests/test_mesa_logging.py

@@ -0,0 +1,47 @@
+"""Unit tests for mesa_logging."""
+
+import logging
+
+import pytest
+
+from mesa import mesa_logging
+
+
+@pytest.fixture
+def tear_down():
+    """Pytest fixture to ensure all logging is disabled after testing."""
+    yield None
+    mesa_logging._logger = None
+    ema_logger = logging.getLogger(mesa_logging.LOGGER_NAME)
+    ema_logger.handlers = []
+
+
+def test_get_logger():
+    """Test get_logger."""
+    mesa_logging._rootlogger = None
+    logger = mesa_logging.get_rootlogger()
+    assert logger == logging.getLogger(mesa_logging.LOGGER_NAME)
+    assert len(logger.handlers) == 1
+    assert isinstance(logger.handlers[0], logging.NullHandler)
+
+    logger = mesa_logging.get_rootlogger()
+    assert logger, logging.getLogger(mesa_logging.LOGGER_NAME)
+    assert len(logger.handlers) == 1
+    assert isinstance(logger.handlers[0], logging.NullHandler)
+
+
+def test_log_to_stderr():
+    """Test log_to_stderr."""
+    mesa_logging._rootlogger = None
+    logger = mesa_logging.log_to_stderr(mesa_logging.DEBUG)
+    assert len(logger.handlers) == 2
+    assert logger.level == mesa_logging.DEBUG
+
+    mesa_logging._rootlogger = None
+    logger = mesa_logging.log_to_stderr()
+    assert len(logger.handlers) == 2
+    assert logger.level == mesa_logging.DEFAULT_LEVEL
+
+    logger = mesa_logging.log_to_stderr()
+    assert len(logger.handlers) == 2
+    assert logger.level == mesa_logging.DEFAULT_LEVEL