Source code for tianshou.utils.logging
"""
Partial copy of sensai.util.logging
"""
# ruff: noqa
import atexit
import logging as lg
import sys
from collections.abc import Callable
from datetime import datetime
from io import StringIO
from logging import *
from typing import Any, TypeVar, cast
log = getLogger(__name__) # type: ignore
LOG_DEFAULT_FORMAT = "%(levelname)-5s %(asctime)-15s %(name)s:%(funcName)s - %(message)s"
# Holds the log format that is configured by the user (using function `configure`), such
# that it can be reused in other places
_logFormat = LOG_DEFAULT_FORMAT
[docs]
def set_numerical_fields_to_precision(data: dict[str, Any], precision: int = 3) -> dict[str, Any]:
"""Returns a copy of the given dictionary with all numerical values rounded to the given precision.
Note: does not recurse into nested dictionaries.
:param data: a dictionary
:param precision: the precision to be used
"""
result = {}
for k, v in data.items():
if isinstance(v, float):
v = round(v, precision)
result[k] = v
return result
[docs]
def remove_log_handlers() -> None:
"""Removes all current log handlers."""
logger = getLogger()
while logger.hasHandlers():
logger.removeHandler(logger.handlers[0])
[docs]
def remove_log_handler(handler: Handler) -> None:
getLogger().removeHandler(handler)
[docs]
def is_log_handler_active(handler: Handler) -> bool:
"""Checks whether the given handler is active.
:param handler: a log handler
:return: True if the handler is active, False otherwise
"""
return handler in getLogger().handlers
# noinspection PyShadowingBuiltins
T = TypeVar("T")
# noinspection PyShadowingBuiltins
[docs]
def run_main(
main_fn: Callable[[], T], format: str = LOG_DEFAULT_FORMAT, level: int = lg.DEBUG
) -> T | None:
"""Configures logging with the given parameters, ensuring that any exceptions that occur during
the execution of the given function are logged.
Logs two additional messages, one before the execution of the function, and one upon its completion.
:param main_fn: the function to be executed
:param format: the log message format
:param level: the minimum log level
:return: the result of `main_fn`
"""
configure(format=format, level=level)
log.info("Starting") # type: ignore
try:
result = main_fn()
log.info("Done") # type: ignore
return result
except Exception as e:
log.error("Exception during script execution", exc_info=e) # type: ignore
return None
[docs]
def run_cli(
main_fn: Callable[..., T], format: str = LOG_DEFAULT_FORMAT, level: int = lg.DEBUG
) -> T | None:
"""
Configures logging with the given parameters and runs the given main function as a
CLI using `jsonargparse` (which is configured to also parse attribute docstrings, such
that dataclasses can be used as function arguments).
Using this function requires that `jsonargparse` and `docstring_parser` be available.
Like `run_main`, two additional log messages will be logged (at the beginning and end
of the execution), and it is ensured that all exceptions will be logged.
:param main_fn: the function to be executed
:param format: the log message format
:param level: the minimum log level
:return: the result of `main_fn`
"""
from jsonargparse import set_docstring_parse_options, CLI
set_docstring_parse_options(attribute_docstrings=True)
return run_main(lambda: CLI(main_fn), format=format, level=level)
[docs]
def datetime_tag() -> str:
""":return: a string tag for use in log file names which contains the current date and time (compact but readable)"""
return datetime.now().strftime("%Y%m%d-%H%M%S")
_fileLoggerPaths: list[str] = []
_isAtExitReportFileLoggerRegistered = False
_memoryLogStream: StringIO | None = None
def _at_exit_report_file_logger() -> None:
for path in _fileLoggerPaths:
print(f"A log file was saved to {path}")
[docs]
def add_file_logger(path: str, register_atexit: bool = True) -> FileHandler:
global _isAtExitReportFileLoggerRegistered
log.info(f"Logging to {path} ...") # type: ignore
handler = FileHandler(path)
handler.setFormatter(Formatter(_logFormat))
Logger.root.addHandler(handler)
_fileLoggerPaths.append(path)
if not _isAtExitReportFileLoggerRegistered and register_atexit:
atexit.register(_at_exit_report_file_logger)
_isAtExitReportFileLoggerRegistered = True
return handler
[docs]
def add_memory_logger() -> None:
"""Enables in-memory logging (if it is not already enabled), i.e. all log statements are written to a memory buffer and can later be
read via function `get_memory_log()`.
"""
global _memoryLogStream
if _memoryLogStream is not None:
return
_memoryLogStream = StringIO()
handler = StreamHandler(_memoryLogStream)
handler.setFormatter(Formatter(_logFormat))
Logger.root.addHandler(handler)
[docs]
def get_memory_log() -> Any:
""":return: the in-memory log (provided that `add_memory_logger` was called beforehand)"""
assert _memoryLogStream is not None, "This should not have happened and might be a bug."
return _memoryLogStream.getvalue()
[docs]
class FileLoggerContext:
def __init__(self, path: str, enabled: bool = True):
self.enabled = enabled
self.path = path
self._log_handler: Handler | None = None
def __enter__(self) -> None:
if self.enabled:
self._log_handler = add_file_logger(self.path, register_atexit=False)
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
if self._log_handler is not None:
remove_log_handler(self._log_handler)