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 """Custom log formatter with ANSI color support by log level.""" 

29 

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

31 

32 LVL_COLOR = { 

33 logging.DEBUG: Bcolors.OKBLUE, 

34 logging.INFO: "", 

35 logging.WARNING: Bcolors.WARNING, 

36 logging.ERROR: Bcolors.FAIL, 

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

38 } 

39 

40 @override 

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

42 """Format a log record with ANSI color based on log level. 

43 

44 Args: 

45 record: The log record to format. 

46 

47 Returns: 

48 Formatted log message string with ANSI color codes. 

49 """ 

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

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

52 formatter = logging.Formatter(log_fmt) 

53 return formatter.format(record) 

54 

55 @classmethod 

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

57 """Initialize logging with the custom color formatter. 

58 

59 Args: 

60 level: Logging level (int or str). 

61 format: Log message format string. 

62 """ 

63 # TODO : forward generic parameters of init_log to basicConfig 

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

65 

66 # We upgrade format for our color console formater 

67 cls.MSG_FMT = format 

68 

69 # Adding color format to stream handler 

70 logger = logging.getLogger() 

71 for handler in logger.handlers: 

72 if isinstance(handler, logging.StreamHandler): 

73 handler.setFormatter(CustomFormatter())