Coverage for credoai/utils/logging.py: 76%

1from collections import deque 

2from io import StringIO 

3from logging import FileHandler, Formatter, Handler, StreamHandler, getLogger 

4from os.path import join 

5from sys import stdout 

6 

7 

8class TailLogHandler(Handler): 

9 def __init__(self, log_queue): 

10 Handler.__init__(self) 

11 self.log_queue = log_queue 

12 

13 def emit(self, record): 

14 self.log_queue.append(self.format(record)) 

15 

16 

17class TailLogger(object): 

18 def __init__(self, maxlen): 

19 self._log_queue = deque(maxlen=maxlen) 

20 self._log_handler = TailLogHandler(self._log_queue) 

21 

22 def contents(self): 

23 return "\n".join(self._log_queue) 

24 

25 @property 

26 def log_handler(self): 

27 return self._log_handler 

28 

29 

30class Logger: 

31 def __init__( 

32 self, name, path=None, record_stream=True, logging_level="info", formatter=None 

33 ): 

34 self.file_path = None 

35 self.stream = None 

36 if formatter is None: 

37 formatter = Formatter( 

38 "%(asctime)s - %(name)s - %(levelname)s - %(message)s" 

39 ) 

40 self.formatter = formatter 

41 self.logger = self.setup_logger(name, logging_level) 

42 self.log_capture_string = StringIO() 

43 if record_stream: 

44 self.stream = self.setup_stream() 

45 if path: 

46 self.file_path = join(path, f"{name}.log") 

47 self.setup_file() 

48 

49 def setup_logger(self, name, logging_level): 

50 logger = getLogger(name) 

51 handler = StreamHandler(stdout) 

52 handler.setFormatter(self.formatter) 

53 logger.addHandler(handler) 

54 logger.setLevel(logging_level) 

55 return logger 

56 

57 def setup_stream(self, tail_length=100): 

58 tail = TailLogger(tail_length) 

59 log_handler = tail.log_handler 

60 log_handler.setFormatter(self.formatter) 

61 self.logger.addHandler(log_handler) 

62 return tail 

63 

64 def setup_file(self): 

65 file_handler = FileHandler(self.file_path) 

66 file_handler.setFormatter(self.formatter) 

67 self.logger.addHandler(file_handler) 

68 

69 

70def setup_logger(name="lens", path=None, record_stream=False, logging_level="INFO"): 

71 tmp = Logger(name, path, record_stream, logging_level) 

72 return tmp.logger, tmp.stream 

73 

74 

75global_logger, global_tail = setup_logger(path=".")