From 935d4a36909006f7f43debd08d1ec44d5dd04771 Mon Sep 17 00:00:00 2001
From: Veniamin Malefioudakis <58257722+benmalef@users.noreply.github.com>
Date: Mon, 15 Jul 2024 13:23:08 +0300
Subject: [PATCH] update gandlf_logger

- Update gandlf_logging_setup
- Update logging_cofing
- Update setup.py
---
 GANDLF/logging_config.yaml    | 10 ++++++--
 GANDLF/utils/gandlf_logger.py | 48 ++++++++++++++++++++++++-----------
 setup.py                      |  3 ++-
 testing/test_full.py          |  6 +++--
 4 files changed, 47 insertions(+), 20 deletions(-)

diff --git a/GANDLF/logging_config.yaml b/GANDLF/logging_config.yaml
index 89b6100b6..6bc1fbdac 100644
--- a/GANDLF/logging_config.yaml
+++ b/GANDLF/logging_config.yaml
@@ -3,8 +3,15 @@ formatters:
   detailed:
     format: "%(asctime)s - %(name)s - %(levelname)s - %(module)s:%(funcName)s:%(lineno)d - %(message)s"
     datefmt: "%Y-%m-%d %H:%M:%S"
+    log_colors:
+      DEBUG: "white"
+      INFO: "green"
+      WARNING: "yellow"
+      ERROR: "red"
+      CRITICAL: "bold_red"
   simple:
-    format: "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    (): colorlog.ColoredFormatter
+    format: "%(log_color)s%(asctime)s - %(name)s - %(levelname)s - %(message)s"
     datefmt: "%Y-%m-%d %H:%M:%S"
 filters:
   warnings_filter:
@@ -28,7 +35,6 @@ handlers:
     class: logging.handlers.RotatingFileHandler
     level: DEBUG
     formatter: detailed
-    filename: tmp/gandlf/gandlf.log
     maxBytes: 10485760
     backupCount: 2
 loggers: # you can add your customized logger
diff --git a/GANDLF/utils/gandlf_logger.py b/GANDLF/utils/gandlf_logger.py
index 2d8412ca4..7f6c309ee 100644
--- a/GANDLF/utils/gandlf_logger.py
+++ b/GANDLF/utils/gandlf_logger.py
@@ -2,31 +2,49 @@
 import yaml
 from pathlib import Path
 from importlib import resources
+import os
+import colorlog
 
 
-def gandlf_logger_setup(config_path="logging_config.yaml") -> logging.Logger:
+def gandlf_logger_setup(log_dir=None, config_path="logging_config.yaml"):
     """
-    It sets up the logger. Read from logging_config.
+    It sets up the logger. Reads from logging_config.
+    If log_dir is None, the logs are flashed to console.
     Args:
-        logger_name (str): logger name, the name should be the same in the logging_config
+        log_dir (str): dir path for saving the logs
         config_path (str): file path for the configuration
-    Returns:
-        logging.Logger
-    """
 
-    # create dir for storing the messages
-    current_dir = Path.cwd()
-    directory = Path.joinpath(current_dir, "tmp/gandlf")
-    directory.mkdir(parents=True, exist_ok=True)
+    """
 
-    with resources.open_text("GANDLF", config_path) as file:
-        config_dict = yaml.safe_load(file)
-        logging.config.dictConfig(config_dict)
+    if log_dir == None:  # flash logs
+        formatter = colorlog.ColoredFormatter(
+            "%(log_color)s%(asctime)s - %(name)s - %(levelname)s - %(module)s:%(funcName)s:%(lineno)d - %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+            log_colors={
+                "DEBUG": "blue",
+                "INFO": "green",
+                "WARNING": "yellow",
+                "ERROR": "red",
+                "CRITICAL": "bold_red",
+            },
+        )
+        console_handler = logging.StreamHandler()
+        console_handler.setFormatter(formatter)
+        logging.root.setLevel(logging.DEBUG)
+        logging.root.addHandler(console_handler)
+
+    else:  # create the log file
+        output_dir = Path(log_dir)
+        Path(output_dir).mkdir(parents=True, exist_ok=True)
+        with resources.open_text("GANDLF", config_path) as file:
+            config_dict = yaml.safe_load(file)
+            config_dict["handlers"]["rotatingFileHandler"]["filename"] = str(
+                Path.joinpath(output_dir, "gandlf.log")
+            )
+            logging.config.dictConfig(config_dict)
 
     logging.captureWarnings(True)
 
-    # return logging.getLogger(logger_name)
-
 
 class InfoOnlyFilter(logging.Filter):
     """
diff --git a/setup.py b/setup.py
index a5f9d5c09..1ee6a3f72 100644
--- a/setup.py
+++ b/setup.py
@@ -31,7 +31,7 @@
 ]
 
 # Any extra files should be located at `GANDLF` module folder (not in repo root)
-extra_files = ["logging_config.yml"]
+extra_files = ["logging_config.yaml"]
 toplevel_package_excludes = ["testing*"]
 
 # specifying version for `black` separately because it is also used to [check for lint](https://github.com/mlcommons/GaNDLF/blob/master/.github/workflows/black.yml)
@@ -80,6 +80,7 @@
     "deprecated",
     "packaging==24.0",
     "typer==0.9.0",
+    "colorlog",
 ]
 
 if __name__ == "__main__":
diff --git a/testing/test_full.py b/testing/test_full.py
index ce0c78293..df080e32b 100644
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -3190,16 +3190,18 @@ def test_generic_data_split():
 def test_generic_logging(capsys):
     print("52: Starting test for logging")
 
-    gandlf_logger_setup()
+    gandlf_logger_setup("testing/log")
     message = "Testing logging"
 
     logging.debug(message)
 
     # tests if the message is in the file.log
-    with open("tmp/gandlf/gandlf.log", "r") as log_file:
+    with open("testing/log/gandlf.log", "r") as log_file:
         logs = log_file.read()
         assert message in logs
 
+    shutil.rmtree("testing/log")
+
     # test the stout info level. The stout must show only INFO messages
     message = "Testing stout logging"
     logging.info(message)