add logger_utils

src/egamma_tnp/utils/logger_utils.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import logging
+
+from rich.console import Console
+from rich.logging import RichHandler
+
+LOGGER_NAME = "egamma_tnp"
+
+
+def setup_logger(level: str = "INFO", logfile: str | None = None, time: bool | None = False) -> logging.Logger:
+    """Setup a logger that uses RichHandler to write the same message both in stdout
+    and in a log file called logfile. Level of information can be customized and
+    dumping a logfile is optional.
+
+    :param level: level of information
+    :type level: str, optional
+    :param logfile: file where information are stored
+    :type logfile: str
+    """
+    logger = logging.getLogger(LOGGER_NAME)  # need to give it a name, otherwise *way* too much info gets printed out from e.g. numba
+
+    # Set up level of information
+    possible_levels = ["INFO", "DEBUG"]
+    if level not in possible_levels:
+        raise ValueError("Passed wrong level for the logger. Allowed levels are: {}".format(", ".join(possible_levels)))
+    logger.setLevel(getattr(logging, level))
+
+    formatter = logging.Formatter("%(message)s")
+    if time:
+        formatter = logging.Formatter("%(asctime)s %(message)s")
+
+    # Set up stream handler (for stdout)
+    stream_handler = RichHandler(show_time=False, rich_tracebacks=True)
+    stream_handler.setFormatter(formatter)
+    logger.addHandler(stream_handler)
+
+    # Set up file handler (for logfile)
+    if logfile:
+        file_handler = RichHandler(
+            show_time=False,
+            rich_tracebacks=True,
+            console=Console(file=open(logfile, "w")),
+        )
+        file_handler.setFormatter(formatter)
+        logger.addHandler(file_handler)
+
+    return logger