Coverage for muutils / logger / simplelogger.py: 52%

1from __future__ import annotations 

2 

3import json 

4import sys 

5import time 

6import typing 

7from typing import Any, TextIO, Union 

8 

9from muutils.json_serialize import JSONitem, json_serialize 

10 

11 

12class NullIO: 

13 """null IO class""" 

14 

15 def __init__(self) -> None: 

16 pass 

17 

18 def write(self, msg: str) -> int: 

19 """write to nothing! this throws away the message""" 

20 return len(msg) 

21 

22 def flush(self) -> None: 

23 """flush nothing! this is a no-op""" 

24 pass 

25 

26 def close(self) -> None: 

27 """close nothing! this is a no-op""" 

28 pass 

29 

30 

31AnyIO = Union[TextIO, NullIO] 

32 

33 

34class SimpleLogger: 

35 """logs training data to a jsonl file""" 

36 

37 def __init__( 

38 self, 

39 log_path: str | None = None, 

40 log_file: AnyIO | None = None, 

41 timestamp: bool = True, 

42 ): 

43 self._timestamp: bool = timestamp 

44 self._log_path: str | None = log_path 

45 

46 self._log_file_handle: AnyIO 

47 

48 if (log_path is None) and (log_file is None): 

49 print( 

50 "[logger_internal] # no log file specified, will only write to console", 

51 sys.stderr, 

52 ) 

53 self._log_file_handle = sys.stdout 

54 

55 elif (log_path is not None) and (log_file is not None): 

56 raise ValueError( 

57 "cannot specify both log_path and log_file, use streams in `SimpleLogger`" 

58 ) 

59 else: 

60 # now exactly one of the two is None 

61 if log_file is not None: 

62 self._log_file_handle = log_file 

63 else: 

64 assert log_path is not None 

65 self._log_file_handle = open(log_path, "w", encoding="utf-8") 

66 

67 def log(self, msg: JSONitem, *, console_print: bool = False, **kwargs: Any) -> None: 

68 """log a message to the log file, and optionally to the console""" 

69 if console_print: 

70 print(msg) 

71 

72 msg_dict: dict[str, Any] 

73 if not isinstance(msg, typing.Mapping): 

74 msg_dict = {"_msg": msg} 

75 else: 

76 msg_dict = dict(typing.cast(typing.Mapping[str, Any], msg)) 

77 

78 if self._timestamp: 

79 msg_dict["_timestamp"] = time.time() 

80 

81 if len(kwargs) > 0: 

82 msg_dict["_kwargs"] = kwargs 

83 

84 self._log_file_handle.write(json.dumps(json_serialize(msg_dict)) + "\n")