Coverage for agentlib/core/logging_.py: 94%

1import logging 

2from typing import TYPE_CHECKING 

3 

4from agentlib.core import environment 

5 

6if TYPE_CHECKING: 

7 from agentlib import Environment 

8 

9 

10class CustomLogger(logging.Logger): 

11 """Subclass of Logger that adds the env_time to the record, allowing it to print 

12 the current time.""" 

13 

14 def __init__(self, name, env: "Environment", level=logging.NOTSET): 

15 super().__init__(name, level) 

16 self.env = env 

17 

18 def makeRecord( 

19 self, 

20 name, 

21 level, 

22 fn, 

23 lno, 

24 msg, 

25 args, 

26 exc_info, 

27 func=None, 

28 extra=None, 

29 sinfo=None, 

30 ): 

31 record = super().makeRecord( 

32 name, level, fn, lno, msg, args, exc_info, func, extra, sinfo 

33 ) 

34 _until = self.env.pretty_until() 

35 _time = self.env.pretty_time() 

36 if _until is environment.UNTIL_UNSET: 

37 # Add "INIT" prefix to clearly indicate initialization phase 

38 record.env_time = f"<INIT>" 

39 elif _until is None: 

40 record.env_time = _time 

41 else: 

42 record.env_time = _time + "/" + _until 

43 return record 

44 

45 

46def create_logger(env: "Environment", name: str) -> CustomLogger: 

47 """Creates a logger that displays the environment time when logging.""" 

48 # Create a custom logger 

49 custom_logger = CustomLogger(name, env=env) 

50 custom_logger.setLevel(logging.root.getEffectiveLevel()) 

51 

52 # Create a formatter 

53 formatter = logging.Formatter("%(env_time)s %(levelname)s: %(name)s: %(message)s") 

54 

55 # Create a StreamHandler and add it to the logger 

56 stream_handler = logging.StreamHandler() 

57 stream_handler.setFormatter(formatter) 

58 custom_logger.addHandler(stream_handler) 

59 

60 # Check if root logger has any FileHandlers and add similar ones to our custom logger 

61 for handler in logging.root.handlers: 

62 if isinstance(handler, logging.FileHandler): 

63 # Create a similar FileHandler for our custom logger 

64 file_handler = logging.FileHandler( 

65 handler.baseFilename, mode=handler.mode, encoding=handler.encoding 

66 ) 

67 file_handler.setFormatter(formatter) 

68 file_handler.setLevel(handler.level) 

69 custom_logger.addHandler(file_handler) 

70 

71 return custom_logger