Devs/rich logging (#105)

FBumann · web-flow · commit 2fed1fd29204 · 2024-12-16T12:39:12.000+01:00
* Using the package rich for formated logging instead of own formater

* Adding logging formater for logging to file

* Re-adding the old custom formater (simplified), to enable colored logs in the python console

* Adding rich.print() to examples

* Changing rich.print() to rich.pprint()

* Adjusted colors of logging to match RichHandler

* Default to not use the rich Handler
diff --git a/examples/Ex00_minmal/minimal_example.py b/examples/Ex00_minmal/minimal_example.py
@@ -1,9 +1,10 @@
 """
 This script shows how to use the flixOpt framework to model a super minimalistic energy system.
 """
+import flixOpt as fx
 
 import numpy as np
-import flixOpt as fx
+from rich.pretty import pprint
 
 if __name__ == '__main__':
 
@@ -54,5 +55,6 @@
     results.plot_operation('District Heating', 'area')
 
     # Print results to the console. Check Results in file or perform more plotting
-    print(calculation.results())
-    print(f'Look into .yaml and .json file for results')
+    pprint(calculation.results())
+    pprint(f'Look into .yaml and .json file for results')
+    pprint(calculation.system_model.main_results)
diff --git a/examples/Ex01_simple/simple_example.py b/examples/Ex01_simple/simple_example.py
@@ -1,8 +1,8 @@
 """
 THis script shows how to use the flixOpt framework to model a simple energy system.
 """
-
 import numpy as np
+from rich.pretty import pprint  # Used for pretty printing
 import flixOpt as fx
 
 if __name__ == '__main__':
@@ -96,4 +96,4 @@
 
     # Convert the results for the storage component to a dataframe and display
     results.to_dataframe('Storage')
-    print(results.all_results)
+    pprint(results.all_results)
diff --git a/examples/Ex02_complex/complex_example.py b/examples/Ex02_complex/complex_example.py
@@ -1,9 +1,9 @@
 """
 This script shows how to use the flixOpt framework to model a more complex energy system.
 """
-
 import numpy as np
 import pandas as pd
+from rich.pretty import pprint  # Used for pretty printing
 import flixOpt as fx
 
 if __name__ == '__main__':
@@ -125,15 +125,15 @@
     flow_system.add_elements(Costs, CO2, PE, Gaskessel, Waermelast, Gasbezug, Stromverkauf, aSpeicher)
     flow_system.add_elements(aKWK2) if use_chp_with_segments else flow_system.add_components(aKWK)
 
-    print(flow_system)  # Get a string representation of the FlowSystem
+    pprint(flow_system)  # Get a string representation of the FlowSystem
 
     # --- Solve FlowSystem ---
     calculation = fx.FullCalculation('Sim1', flow_system, 'pyomo', time_indices)
     calculation.do_modeling()
 
     # Show variables as str (else, you can find them in the results.yaml file
-    print(calculation.system_model.description_of_constraints())
-    print(calculation.system_model.description_of_variables())
+    pprint(calculation.system_model.description_of_constraints())
+    pprint(calculation.system_model.description_of_variables())
 
     calculation.solve(fx.solvers.HighsSolver(mip_gap=0.005, time_limit_seconds=30),  # Specify which solver you want to use and specify parameters
                       save_results='results')  # If and where to save results
diff --git a/examples/Ex03_calculation_types/example_calculation_types.py b/examples/Ex03_calculation_types/example_calculation_types.py
@@ -8,6 +8,7 @@
 
 import numpy as np
 import pandas as pd
+from rich.pretty import pprint  # Used for pretty printing
 
 import flixOpt as fx
 
@@ -148,7 +149,7 @@
         calculation.solve(fx.solvers.HighsSolver())
         calculations['Aggregated'] = calculation
         results['Aggregated'] = calculations['Aggregated'].results()
-
+    pprint(results)
 
     def extract_result(results_data: dict[str, dict], keys: List[str]) -> Dict[str,Union[int, float, np.ndarray]]:
         """
diff --git a/flixOpt/__init__.py b/flixOpt/__init__.py
@@ -3,4 +3,4 @@
 """
 
 from .commons import *
-setup_logging('INFO')
+setup_logging('INFO', use_rich_handler=False)
diff --git a/flixOpt/core.py b/flixOpt/core.py
@@ -8,6 +8,8 @@
 import inspect
 
 import numpy as np
+from rich.logging import RichHandler
+from rich.console import Console
 
 from . import utils
 
@@ -286,57 +288,95 @@ def as_effect_dict_with_ts(name_of_param: str,
     return effect_ts_dict
 
 
-# TODO: Move logging to utils.py
-class CustomFormatter(logging.Formatter):
+class MultilineFormater(logging.Formatter):
+
+    def format(self, record):
+        message_lines = record.getMessage().split('\n')
+
+        # Prepare the log prefix (timestamp + log level)
+        timestamp = self.formatTime(record, self.datefmt)
+        log_level = record.levelname.ljust(8)  # Align log levels for consistency
+        log_prefix = f"{timestamp} | {log_level} |"
+
+        # Format all lines
+        first_line = [f'{log_prefix} {message_lines[0]}']
+        if len(message_lines) > 1:
+            lines = first_line + [f"{log_prefix} {line}" for line in message_lines[1:]]
+        else:
+            lines = first_line
+
+        return '\n'.join(lines)
+
+
+class ColoredMultilineFormater(MultilineFormater):
     # ANSI escape codes for colors
     COLORS = {
-        'DEBUG': '\033[96m',    # Cyan
-        'INFO': '\033[92m',     # Green
-        'WARNING': '\033[93m',  # Yellow
-        'ERROR': '\033[91m',    # Red
-        'CRITICAL': '\033[91m\033[1m',  # Bold Red
+        'DEBUG': '\033[32m',  # Green
+        'INFO': '\033[34m',  # Blue
+        'WARNING': '\033[33m',  # Yellow
+        'ERROR': '\033[31m',  # Red
+        'CRITICAL': '\033[1m\033[31m',  # Bold Red
     }
     RESET = '\033[0m'
 
     def format(self, record):
+        lines = super().format(record).splitlines()
         log_color = self.COLORS.get(record.levelname, self.RESET)
-        original_message = record.getMessage()
-        message_lines = original_message.split('\n')
 
         # Create a formatted message for each line separately
         formatted_lines = []
-        for line in message_lines:
-            temp_record = logging.LogRecord(
-                record.name, record.levelno, record.pathname, record.lineno,
-                line, record.args, record.exc_info, record.funcName, record.stack_info
-            )
-            formatted_line = super().format(temp_record)
-            formatted_lines.append(f"{log_color}{formatted_line}{self.RESET}")
-
-        formatted_message = '\n'.join(formatted_lines)
-        return formatted_message
-
-
-def setup_logging(level_name: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']):
+        for line in lines:
+            formatted_lines.append(f"{log_color}{line}{self.RESET}")
+
+        return '\n'.join(formatted_lines)
+
+
+def _get_logging_handler(log_file: Optional[str] = None,
+                         use_rich_handler: bool = False) -> logging.Handler:
+    """Returns a logging handler for the given log file."""
+    if use_rich_handler and log_file is None:
+        # RichHandler for console output
+        console = Console(width=120)
+        rich_handler = RichHandler(
+            console=console,
+            rich_tracebacks=True,
+            omit_repeated_times=True,
+            show_path=False,
+            log_time_format="%Y-%m-%d %H:%M:%S",
+        )
+        rich_handler.setFormatter(logging.Formatter("%(message)s"))  # Simplified formatting
+
+        return rich_handler
+    elif log_file is None:
+        # Regular Logger with custom formating enabled
+        file_handler = logging.StreamHandler()
+        file_handler.setFormatter(ColoredMultilineFormater(
+            fmt="%(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        ))
+        return file_handler
+    else:
+        # FileHandler for file output
+        file_handler = logging.FileHandler(log_file)
+        file_handler.setFormatter(MultilineFormater(
+            fmt="%(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        ))
+        return file_handler
+
+def setup_logging(default_level: Literal['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'] = 'INFO',
+                  log_file: Optional[str] = 'flixOpt.log',
+                  use_rich_handler: bool = False):
     """Setup logging configuration"""
     logger = logging.getLogger('flixOpt')  # Use a specific logger name for your package
-    logging_level = get_logging_level_by_name(level_name)
-    logger.setLevel(logging_level)
-
+    logger.setLevel(get_logging_level_by_name(default_level))
     # Clear existing handlers
     if logger.hasHandlers():
         logger.handlers.clear()
 
-    # Create console handler
-    c_handler = logging.StreamHandler()
-    c_handler.setLevel(logging_level)
-
-    # Create a clean and aligned formatter
-    log_format = '%(asctime)s - %(levelname)-8s : %(message)s'
-    date_format = '%Y-%m-%d %H:%M:%S'
-    c_format = CustomFormatter(log_format, datefmt=date_format)
-    c_handler.setFormatter(c_format)
-    logger.addHandler(c_handler)
+    logger.addHandler(_get_logging_handler(use_rich_handler=use_rich_handler))
+    if log_file is not None:
+        logger.addHandler(_get_logging_handler(log_file, use_rich_handler=False))
 
     return logger
 
diff --git a/pyproject.toml b/pyproject.toml
@@ -33,6 +33,7 @@ dependencies = [
     "numpy >= 1.21.5, < 2",
     "PyYAML >= 6.0",
     "Pyomo >= 6.4.2",
+    "rich >= 13.0.1",
     "tsam >= 2.3.1",  # Used for time series aggregation
     "highspy >= 1.5.3",  # Default solver
     "pandas >= 2, < 3",  # Used in post-processing
