Coverage for packages / dqm-ml / src / dqm_ml / cli_tools.py: 100%

1"""CLI utilities for DQM-ML v2. 

2 

3This module provides utilities for colored console output and custom 

4logging formatters. 

5""" 

6 

7import logging 

8 

9from typing_extensions import override 

10 

11 

12class Bcolors: 

13 """ANSI color codes for terminal output formatting.""" 

14 

15 HEADER = "\033[95m" 

16 OKBLUE = "\033[94m" 

17 OKCYAN = "\033[96m" 

18 OKGREEN = "\033[92m" 

19 WARNING = "\033[93m" 

20 FAIL = "\033[91m" 

21 ENDC = "\033[0m" 

22 BOLD = "\033[1m" 

23 UNDERLINE = "\033[4m" 

24 

25 

26# TODO move logging in a dedicated module 

27class CustomFormatter(logging.Formatter): 

28 MSG_FMT = "%(asctime)s - %(name)s - %(levelname)s - %(message)s (%(filename)s:%(lineno)d)" 

29 

30 LVL_COLOR = { 

31 logging.DEBUG: Bcolors.OKBLUE, 

32 logging.INFO: "", 

33 logging.WARNING: Bcolors.WARNING, 

34 logging.ERROR: Bcolors.FAIL, 

35 logging.CRITICAL: Bcolors.FAIL + Bcolors.BOLD, 

36 } 

37 

38 @override 

39 def format(self, record: logging.LogRecord) -> str: 

40 color = self.LVL_COLOR.get(record.levelno, "") 

41 log_fmt = color + self.MSG_FMT + Bcolors.ENDC 

42 formatter = logging.Formatter(log_fmt) 

43 return formatter.format(record) 

44 

45 @classmethod 

46 def init_log(cls, level: int | str, format: str = MSG_FMT) -> None: 

47 # TODO : forward generic parameters of init_log to basicConfig 

48 logging.basicConfig(format=format, level=level) 

49 

50 # We upgrade format for our color console formater 

51 cls.MSG_FMT = format 

52 

53 # Adding color format to stream handler 

54 logger = logging.getLogger() 

55 for handler in logger.handlers: 

56 if isinstance(handler, logging.StreamHandler): 

57 handler.setFormatter(CustomFormatter())