Coverage for instanovo/utils/colorlogging.py: 93%

1import logging 

2import os 

3from typing import Literal 

4 

5from rich.console import Console 

6from rich.logging import RichHandler 

7 

8from instanovo.__init__ import get_rank 

9from instanovo.constants import LOGGING_SHOW_PATH, LOGGING_SHOW_TIME, USE_RICH_HANDLER 

10 

11 

12class DynamicRankFormatter(logging.Formatter): 

13 """A formatter that dynamically includes rank information in the logger name.""" 

14 

15 def __init__( 

16 self, 

17 fmt: str | None = None, 

18 datefmt: str | None = None, 

19 style: Literal["%", "{", "$"] = "%", 

20 ) -> None: 

21 super().__init__(fmt, datefmt, style) 

22 

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

24 """Format a log record. 

25 

26 Args: 

27 record (logging.LogRecord): The log record to format. 

28 

29 Returns: 

30 str: The formatted log record 

31 """ 

32 rank = get_rank() 

33 

34 # If we have a rank, append it to the logger name 

35 if rank is not None: 

36 record.msg = f"[RANK {rank}] {record.msg}" 

37 

38 return super().format(record) 

39 

40 

41class ColorLog: 

42 """A logging utility class that integrates with Rich for enhanced console output. 

43 

44 (based on https://stackoverflow.com/a/79225597) 

45 """ 

46 

47 def __init__(self, console: Console, name: str) -> None: 

48 message_format = "%(message)s" # Include logger name in format 

49 

50 aichor_enabled = "AICHOR_LOGS_PATH" in os.environ 

51 

52 rich_handler = RichHandler( 

53 console=console, 

54 show_time=LOGGING_SHOW_TIME and not aichor_enabled, 

55 show_path=LOGGING_SHOW_PATH and not aichor_enabled, 

56 ) 

57 rich_handler.setLevel(logging.INFO) 

58 rich_handler.setFormatter(DynamicRankFormatter(message_format)) 

59 

60 if USE_RICH_HANDLER: 

61 logging.basicConfig( 

62 level=logging.NOTSET, 

63 format=message_format, 

64 datefmt="[%X]", 

65 handlers=[rich_handler], 

66 ) 

67 else: 

68 logging.basicConfig( 

69 level=logging.INFO, 

70 format=message_format, 

71 ) 

72 

73 # Suppress INFO logs from the datasets package 

74 logging.getLogger("datasets").setLevel(logging.ERROR) 

75 

76 self.logger = logging.getLogger(name)