Merge pull request #796 from amas0/logger-updates

Enable/disable logging helpers

Author: WardBrian

cmdstanpy/__init__.py

@@ -42,6 +42,8 @@ def _cleanup_tmpdir() -> None:
     set_make_env,
     show_versions,
     write_stan_json,
+    enable_logging,
+    disable_logging,
 )
 
 __all__ = [
@@ -63,4 +65,6 @@ def _cleanup_tmpdir() -> None:
     'show_versions',
     'rebuild_cmdstan',
     'cmdstan_version',
+    "enable_logging",
+    "disable_logging",
 ]

cmdstanpy/utils/__init__.py

@@ -28,7 +28,7 @@
     windows_short_path,
 )
 from .json import write_stan_json
-from .logging import get_logger
+from .logging import get_logger, enable_logging, disable_logging
 from .stancsv import (
     check_sampler_csv,
     parse_rdump_value,
@@ -144,4 +144,6 @@ def show_versions(output: bool = True) -> str:
     'windows_short_path',
     'wrap_url_progress_hook',
     'write_stan_json',
+    'enable_logging',
+    'disable_logging',
 ]

cmdstanpy/utils/logging.py

@@ -1,25 +1,62 @@
 """
 CmdStanPy logging
 """
+
 import functools
 import logging
+import types
+from contextlib import AbstractContextManager
+from typing import Optional, Type
 
 
 @functools.lru_cache(maxsize=None)
 def get_logger() -> logging.Logger:
     """cmdstanpy logger"""
-    logger = logging.getLogger('cmdstanpy')
-    if len(logger.handlers) == 0:
+    logger = logging.getLogger("cmdstanpy")
+    if not logger.hasHandlers():
         # send all messages to handlers
         logger.setLevel(logging.DEBUG)
         # add a default handler to the logger to INFO and higher
         handler = logging.StreamHandler()
         handler.setLevel(logging.INFO)
         handler.setFormatter(
             logging.Formatter(
-                '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+                "%(asctime)s - %(name)s - %(levelname)s - %(message)s",
                 "%H:%M:%S",
             )
         )
         logger.addHandler(handler)
     return logger
+
+
+class ToggleLogging(AbstractContextManager):
+    def __init__(self, disable: bool) -> None:
+        self.disable = disable
+        self.logger = get_logger()
+        self.prev_state = self.logger.disabled
+        self.logger.disabled = self.disable
+
+    def __repr__(self) -> str:
+        return ""
+
+    def __enter__(self) -> "ToggleLogging":
+        self.logger.disabled = self.disable
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_value: Optional[BaseException],
+        traceback: Optional[types.TracebackType],
+    ) -> None:
+        self.logger.disabled = self.prev_state
+
+
+def enable_logging() -> ToggleLogging:
+    """Enable cmdstanpy logging. Can be used as a context manager"""
+    return ToggleLogging(disable=False)
+
+
+def disable_logging() -> ToggleLogging:
+    """Disable cmdstanpy logging. Can be used as a context manager"""
+    return ToggleLogging(disable=True)

docsrc/api.rst

@@ -139,3 +139,13 @@ write_stan_json
 ===============
 
 .. autofunction:: cmdstanpy.write_stan_json
+
+enable_logging
+===============
+
+.. autofunction:: cmdstanpy.enable_logging
+
+disable_logging
+===============
+
+.. autofunction:: cmdstanpy.disable_logging

docsrc/users-guide/outputs.rst

@@ -60,24 +60,47 @@ You may notice CmdStanPy can produce a lot of output when it is running:
 
     fit = model.sample(data=data_file, show_progress=False)
 
-This output is managed through the built-in :mod:`logging` module. For example, it can be disabled entirely:
+This output is managed through the built-in :mod:`logging` module. For
+convenience, CmdStanPy provides the module-level :meth:`cmdstanpy.enable_logging()`
+and :meth:`cmdstanpy.disable_logging()` functions to simplify logging management.
+
+For example, it can be disabled entirely:
 
 .. ipython:: python
 
-    import logging
-    cmdstanpy_logger = logging.getLogger("cmdstanpy")
-    cmdstanpy_logger.disabled = True
+    import cmdstanpy
+    cmdstanpy.disable_logging()
     # look, no output!
     fit = model.sample(data=data_file, show_progress=False)
 
-Or one can remove the logging handler that CmdStanPy installs by default and install their own for more
-fine-grained control. For example, the following code sends all logs (including the ``DEBUG`` logs, which are hidden by default),
-to a file.
+We can re-enable this by calling ``enable_logging()``:
+
+.. ipython:: python
+
+    cmdstanpy.enable_logging()
+    fit = model.sample(data=data_file, show_progress=False)
+
+
+These functions also work as context managers for more local control:
+
+.. ipython:: python
+
+    with cmdstanpy.disable_logging():
+        fit = model.sample(data=data_file, show_progress=False)
+
+
+For more fine-grained control, one can interact with the underlying ``logging``
+library directly. For example, the following code installs a custom handler
+that sends all logs (including the ``DEBUG`` logs, which are hidden by
+default), to a file.
 
 DEBUG logging is useful primarily to developers or when trying to hunt down an issue.
 
 .. ipython:: python
 
+    import logging
+    cmdstanpy_logger = logging.getLogger("cmdstanpy")
+
     cmdstanpy_logger.disabled = False
     # remove all existing handlers
     cmdstanpy_logger.handlers = []
@@ -93,6 +116,7 @@ DEBUG logging is useful primarily to developers or when trying to hunt down an i
     )
     cmdstanpy_logger.addHandler(handler)
 
+
 Now, if we run the model and check the contents of the file, we will see all the possible logging.
 
 .. ipython:: python

test/test_logging.py

@@ -0,0 +1,67 @@
+"""Logging control tests"""
+
+import logging
+import cmdstanpy
+
+
+def test_disable_logging(caplog):
+    logger = cmdstanpy.utils.logging.get_logger()
+
+    with caplog.at_level(logging.INFO, logger="cmdstanpy"):
+        logger.info("before")
+    assert any("before" in m for m in caplog.messages)
+
+    caplog.clear()
+    cmdstanpy.disable_logging()
+
+    with caplog.at_level(logging.INFO, logger="cmdstanpy"):
+        logger.info("after")
+
+    assert not caplog.messages
+    logger.disabled = False
+
+
+def test_disable_logging_context_manager(caplog):
+    logger = cmdstanpy.utils.logging.get_logger()
+
+    with caplog.at_level(logging.INFO, logger="cmdstanpy"):
+        logger.info("before")
+    assert any("before" in m for m in caplog.messages)
+
+    caplog.clear()
+    with cmdstanpy.disable_logging():
+        with caplog.at_level(logging.INFO, logger="cmdstanpy"):
+            logger.info("inside context manager")
+
+    assert not caplog.messages
+
+    with caplog.at_level(logging.INFO, logger="cmdstanpy"):
+        logger.info("after")
+
+    assert any("after" in m for m in caplog.messages)
+    logger.disabled = False
+
+
+def test_disable_logging_context_manager_nested(caplog):
+    logger = cmdstanpy.utils.logging.get_logger()
+
+    with caplog.at_level(logging.INFO, logger="cmdstanpy"):
+        logger.info("before")
+    assert any("before" in m for m in caplog.messages)
+
+    caplog.clear()
+    with cmdstanpy.disable_logging():
+        with cmdstanpy.enable_logging():
+            with caplog.at_level(logging.INFO, logger="cmdstanpy"):
+                logger.info("inside context manager")
+
+    assert any("inside context manager" in m for m in caplog.messages)
+
+    caplog.clear()
+    with cmdstanpy.enable_logging():
+        with cmdstanpy.disable_logging():
+            with caplog.at_level(logging.INFO, logger="cmdstanpy"):
+                logger.info("inside context manager")
+
+    assert not caplog.messages
+    logger.disabled = False