Coverage for stlog/handler.py: 81%

1from __future__ import annotations 

2 

3import logging 

4import sys 

5 

6from stlog.base import ( 

7 rich_dump_exception_on_console, 

8) 

9from stlog.formatter import HumanFormatter, RichHumanFormatter 

10 

11RICH_INSTALLED: bool = False 

12try: 

13 from rich.console import Console 

14 

15 RICH_INSTALLED = True 

16except ImportError: 

17 pass 

18 

19 

20class CustomRichHandler(logging.StreamHandler): 

21 """A custom StreamHandler that uses Rich Console to print log records. 

22 

23 If Rich is not installed or if the stream is not a terminal, 

24 this handler will fallback to the default StreamHandler. 

25 

26 """ 

27 

28 def __init__(self, stream=None, **kwargs): 

29 super().__init__(stream=stream) 

30 self.console = None 

31 self.force_terminal = kwargs.get("force_terminal", False) 

32 if RICH_INSTALLED: 

33 # => let's make a rich console 

34 self.console = Console( 

35 file=stream if stream else sys.stderr, 

36 force_terminal=True if self.force_terminal else None, 

37 highlight=False, 

38 ) 

39 if not (self.force_terminal or self.console.is_terminal): 

40 # Let's fallback to a basic stream handler 

41 self.console = None 

42 if self.console is None: 

43 # Let's default to a human formatter 

44 self.formatter = HumanFormatter() 

45 else: 

46 # Let's use a rich formatter 

47 self.formatter = RichHumanFormatter() 

48 

49 def _rich_emit(self, record: logging.LogRecord): 

50 assert self.console is not None 

51 assert self.formatter is not None 

52 self.console.print(self.formatter.format(record)) 

53 if record.exc_info: 

54 exc_type, exc_value, exc_traceback = record.exc_info 

55 assert exc_type is not None 

56 assert exc_value is not None 

57 rich_dump_exception_on_console( 

58 self.console, exc_type, exc_value, exc_traceback 

59 ) 

60 

61 def emit(self, record: logging.LogRecord): 

62 if self.console is None or self.formatter is None: 

63 return super().emit(record) 

64 else: 

65 return self._rich_emit(record)