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

1import logging 

2from typing import TYPE_CHECKING 

3 

4if TYPE_CHECKING: 

5 from agentlib import Environment 

6 

7 

8class CustomLogger(logging.Logger): 

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

10 the current time.""" 

11 

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

13 super().__init__(name, level) 

14 self.env = env 

15 

16 def makeRecord( 

17 self, 

18 name, 

19 level, 

20 fn, 

21 lno, 

22 msg, 

23 args, 

24 exc_info, 

25 func=None, 

26 extra=None, 

27 sinfo=None, 

28 ): 

29 record = super().makeRecord( 

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

31 ) 

32 record.env_time = self.env.pretty_time() 

33 return record 

34 

35 

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

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

38 # Create a custom logger 

39 custom_logger = CustomLogger(name, env=env) 

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

41 

42 # Create a formatter 

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

44 

45 # Create a StreamHandler and add it to the logger 

46 stream_handler = logging.StreamHandler() 

47 stream_handler.setFormatter(formatter) 

48 custom_logger.addHandler(stream_handler) 

49 return custom_logger