41class Logger(SimpleLogger):
 42    """logger with more features, including log levels and streams
 43
 44    # Parameters:
 45            - `log_path : str | None`
 46            default log file path
 47            (defaults to `None`)
 48            - `log_file : AnyIO | None`
 49            default log io, should have a `.write()` method (pass only this or `log_path`, not both)
 50            (defaults to `None`)
 51            - `timestamp : bool`
 52            whether to add timestamps to every log message (under the `_timestamp` key)
 53            (defaults to `True`)
 54            - `default_level : int`
 55            default log level for streams/messages that don't specify a level
 56            (defaults to `0`)
 57            - `console_print_threshold : int`
 58            log level at which to print to the console, anything greater will not be printed unless overridden by `console_print`
 59            (defaults to `50`)
 60            - `level_header : HeaderFunction`
 61            function for formatting log messages when printing to console
 62            (defaults to `HEADER_FUNCTIONS["md"]`)
 63    - `keep_last_msg_time : bool`
 64            whether to keep the last message time
 65            (defaults to `True`)
 66
 67
 68    # Raises:
 69            - `ValueError` : _description_
 70    """
 71
 72    def __init__(
 73        self,
 74        log_path: str | None = None,
 75        log_file: AnyIO | None = None,
 76        default_level: int = 0,
 77        console_print_threshold: int = 50,
 78        level_header: HeaderFunction = HEADER_FUNCTIONS["md"],
 79        streams: dict[str | None, LoggingStream] | Sequence[LoggingStream] = (),
 80        keep_last_msg_time: bool = True,
 81        # junk args
 82        timestamp: bool = True,
 83        **kwargs: Any,
 84    ) -> None:
 85        # junk arg checking
 86        # ==================================================
 87        if len(kwargs) > 0:
 88            raise ValueError(f"unrecognized kwargs: {kwargs}")
 89
 90        if not timestamp:
 91            raise ValueError(
 92                "timestamp must be True -- why would you not want timestamps?"
 93            )
 94
 95        # timing
 96        # ==================================================
 97        # timing compares
 98        self._keep_last_msg_time: bool = keep_last_msg_time
 99        # TODO: handle per stream?
100        self._last_msg_time: float | None = time.time()
101
102        # basic setup
103        # ==================================================
104        # init BaseLogger
105        super().__init__(log_file=log_file, log_path=log_path, timestamp=timestamp)
106
107        # level-related
108        self._console_print_threshold: int = console_print_threshold
109        self._default_level: int = default_level
110
111        # set up streams
112        self._streams: dict[str | None, LoggingStream] = (
113            streams
114            if isinstance(streams, typing.Mapping)
115            else {s.name: s for s in streams}
116        )
117        # default error stream
118        if "error" not in self._streams:
119            self._streams["error"] = LoggingStream(
120                "error",
121                aliases={
122                    "err",
123                    "except",
124                    "Exception",
125                    "exception",
126                    "exceptions",
127                    "errors",
128                },
129            )
130
131        # check alias duplicates
132        alias_set: set[str | None] = set()
133        for stream in self._streams.values():
134            for alias in stream.aliases:
135                if alias in alias_set:
136                    raise ValueError(f"alias {alias} is already in use")
137                alias_set.add(alias)
138
139        # add aliases
140        for stream in tuple(self._streams.values()):
141            for alias in stream.aliases:
142                if alias not in self._streams:
143                    self._streams[alias] = stream
144
145        # print formatting
146        self._level_header: HeaderFunction = level_header
147
148        print({k: str(v) for k, v in self._streams.items()})
149
150    def _exception_context(
151        self,
152        stream: str = "error",
153        # level: int = -256,
154        # **kwargs,
155    ) -> ExceptionContext:
156        import sys
157
158        s: LoggingStream = self._streams[stream]
159        handler = s.handler if s.handler is not None else sys.stderr
160        return ExceptionContext(stream=handler)
161
162    def log(
163        self,
164        msg: JSONitem = None,
165        *,
166        lvl: int | None = None,
167        stream: str | None = None,
168        console_print: bool = False,
169        extra_indent: str = "",
170        **kwargs: Any,
171    ) -> None:
172        """logging function
173
174        ### Parameters:
175         - `msg : JSONitem`
176           message (usually string or dict) to be logged
177         - `lvl : int | None`
178           level of message (lower levels are more important)
179           (defaults to `None`)
180         - `console_print : bool`
181           override `console_print_threshold` setting
182           (defaults to `False`)
183         - `stream : str | None`
184           whether to log to a stream (defaults to `None`), which logs to the default `None` stream
185           (defaults to `None`)
186        """
187
188        # add to known stream names if not present
189        if stream not in self._streams:
190            self._streams[stream] = LoggingStream(stream)
191
192        # set default level to either global or stream-specific default level
193        # ========================================
194        if lvl is None:
195            if stream is None:
196                lvl = self._default_level
197            else:
198                if self._streams[stream].default_level is not None:
199                    lvl = self._streams[stream].default_level
200                else:
201                    lvl = self._default_level
202
203        assert lvl is not None, "lvl should not be None at this point"
204
205        # print to console with formatting
206        # ========================================
207        _printed: bool = False
208        if console_print or (lvl <= self._console_print_threshold):
209            # add some formatting
210            print(
211                self._level_header(
212                    msg=msg,
213                    lvl=lvl,
214                    stream=stream,
215                    extra_indent=extra_indent,
216                )
217            )
218
219            # store the last message time
220            if self._last_msg_time is not None:
221                self._last_msg_time = time.time()
222
223            _printed = True
224
225        # convert and add data
226        # ========================================
227        # converting to dict
228        msg_dict: dict[str, Any]
229        if not isinstance(msg, typing.Mapping):
230            msg_dict = {"_msg": msg}
231        else:
232            msg_dict = dict(typing.cast(typing.Mapping[str, Any], msg))
233
234        # level+stream metadata
235        if lvl is not None:
236            msg_dict["_lvl"] = lvl
237
238        # msg_dict["_stream"] = stream # moved to LoggingStream
239
240        # extra data in kwargs
241        if len(kwargs) > 0:
242            msg_dict["_kwargs"] = kwargs
243
244        # add default contents (timing, etc)
245        msg_dict = {
246            **{k: v() for k, v in self._streams[stream].default_contents.items()},
247            **msg_dict,
248        }
249
250        # write
251        # ========================================
252        logfile_msg: str = json.dumps(json_serialize(msg_dict)) + "\n"
253        if (
254            (stream is None)
255            or (stream not in self._streams)
256            or (self._streams[stream].handler is None)
257        ):
258            # write to the main log file if no stream is specified
259            self._log_file_handle.write(logfile_msg)
260        else:
261            # otherwise, write to the stream-specific file
262            s_handler: AnyIO | None = self._streams[stream].handler
263            if s_handler is not None:
264                s_handler.write(logfile_msg)
265            else:
266                raise ValueError(
267                    f"stream handler is None! something in the logging stream setup is wrong:\n{self}"
268                )
269
270        # if it was important enough to print, flush all streams
271        if _printed:
272            self.flush_all()
273
274    def log_elapsed_last(
275        self,
276        lvl: int | None = None,
277        stream: str | None = None,
278        console_print: bool = True,
279        **kwargs: Any,
280    ) -> None:
281        """logs the time elapsed since the last message was printed to the console (in any stream)"""
282        if self._last_msg_time is None:
283            raise ValueError("no last message time!")
284        else:
285            self.log(
286                {"elapsed_time": round(time.time() - self._last_msg_time, 6)},
287                lvl=(lvl if lvl is not None else self._console_print_threshold),
288                stream=stream,
289                console_print=console_print,
290                **kwargs,
291            )
292
293    def flush_all(self):
294        """flush all streams"""
295
296        self._log_file_handle.flush()
297
298        for stream in self._streams.values():
299            if stream.handler is not None:
300                stream.handler.flush()
301
302    def __getattr__(self, stream: str) -> Callable[..., Any]:
303        if stream.startswith("_"):
304            raise AttributeError(f"invalid stream name {stream} (no underscores)")
305        return partial(self.log, stream=stream)
306
307    def __getitem__(self, stream: str) -> Callable[..., Any]:
308        return partial(self.log, stream=stream)
309
310    def __call__(self, *args: Any, **kwargs: Any) -> None:
311        self.log(*args, **kwargs)

 18@dataclass
 19class LoggingStream:
 20    """properties of a logging stream
 21
 22    - `name: str` name of the stream
 23    - `aliases: set[str]` aliases for the stream
 24            (calls to these names will be redirected to this stream. duplicate alises will result in errors)
 25            TODO: perhaps duplicate alises should result in duplicate writes?
 26    - `file: str|bool|AnyIO|None` file to write to
 27            - if `None`, will write to standard log
 28            - if `True`, will write to `name + ".log"`
 29            - if `False` will "write" to `NullIO` (throw it away)
 30            - if a string, will write to that file
 31            - if a fileIO type object, will write to that object
 32    - `default_level: int|None` default level for this stream
 33    - `default_contents: dict[str, Callable[[], Any]]` default contents for this stream
 34    - `last_msg: tuple[float, Any]|None` last message written to this stream (timestamp, message)
 35    """
 36
 37    name: str | None
 38    aliases: set[str | None] = field(default_factory=set)
 39    file: str | bool | AnyIO | None = None
 40    default_level: int | None = None
 41    default_contents: dict[str, Callable[[], Any]] = field(default_factory=dict)
 42    handler: AnyIO | None = None
 43
 44    # TODO: implement last-message caching
 45    # last_msg: tuple[float, Any]|None = None
 46
 47    def make_handler(self) -> AnyIO | None:
 48        if self.file is None:
 49            return None
 50        elif isinstance(self.file, str):
 51            # if its a string, open a file
 52            return open(
 53                self.file,
 54                "w",
 55                encoding="utf-8",
 56            )
 57        elif isinstance(self.file, bool):
 58            # if its a bool and true, open a file with the same name as the stream (in the current dir)
 59            # TODO: make this happen in the same dir as the main logfile?
 60            if self.file:
 61                return open(  # type: ignore[return-value]
 62                    f"{sanitize_fname(self.name)}.log.jsonl",
 63                    "w",
 64                    encoding="utf-8",
 65                )
 66            else:
 67                return NullIO()
 68        else:
 69            # if its neither, check it has `.write()` and `.flush()` methods
 70            if (
 71                (
 72                    not hasattr(self.file, "write")
 73                    or (not callable(self.file.write))
 74                    or (not hasattr(self.file, "flush"))
 75                    or (not callable(self.file.flush))
 76                )
 77                or (not hasattr(self.file, "close"))
 78                or (not callable(self.file.close))
 79            ):
 80                raise ValueError(f"stream {self.name} has invalid handler {self.file}")
 81            # ignore type check because we know it has a .write() method,
 82            # assume the user knows what they're doing
 83            return self.file  # type: ignore
 84
 85    def __post_init__(self):
 86        self.aliases = set(self.aliases)
 87        if any(x.startswith("_") for x in self.aliases if x is not None):
 88            raise ValueError(
 89                "stream names or aliases cannot start with an underscore, sorry"
 90            )
 91        self.aliases.add(self.name)
 92        self.default_contents["_timestamp"] = time.time
 93        self.default_contents["_stream"] = lambda: self.name
 94        self.handler = self.make_handler()
 95
 96    def __del__(self):
 97        if self.handler is not None:
 98            self.handler.flush()
 99            self.handler.close()
100
101    @override
102    def __str__(self):
103        return f"LoggingStream(name={self.name}, aliases={self.aliases}, file={self.file}, default_level={self.default_level}, default_contents={self.default_contents})"

35class SimpleLogger:
36    """logs training data to a jsonl file"""
37
38    def __init__(
39        self,
40        log_path: str | None = None,
41        log_file: AnyIO | None = None,
42        timestamp: bool = True,
43    ):
44        self._timestamp: bool = timestamp
45        self._log_path: str | None = log_path
46
47        self._log_file_handle: AnyIO
48
49        if (log_path is None) and (log_file is None):
50            print(
51                "[logger_internal] # no log file specified, will only write to console",
52                sys.stderr,
53            )
54            self._log_file_handle = sys.stdout
55
56        elif (log_path is not None) and (log_file is not None):
57            raise ValueError(
58                "cannot specify both log_path and log_file, use streams in `SimpleLogger`"
59            )
60        else:
61            # now exactly one of the two is None
62            if log_file is not None:
63                self._log_file_handle = log_file
64            else:
65                assert log_path is not None
66                self._log_file_handle = open(log_path, "w", encoding="utf-8")
67
68    def log(self, msg: JSONitem, *, console_print: bool = False, **kwargs: Any) -> None:
69        """log a message to the log file, and optionally to the console"""
70        if console_print:
71            print(msg)
72
73        msg_dict: dict[str, Any]
74        if not isinstance(msg, typing.Mapping):
75            msg_dict = {"_msg": msg}
76        else:
77            msg_dict = dict(typing.cast(typing.Mapping[str, Any], msg))
78
79        if self._timestamp:
80            msg_dict["_timestamp"] = time.time()
81
82        if len(kwargs) > 0:
83            msg_dict["_kwargs"] = kwargs
84
85        self._log_file_handle.write(json.dumps(json_serialize(msg_dict)) + "\n")

 9class TimerContext:
10    """context manager for timing code"""
11
12    def __init__(self) -> None:
13        self.start_time: float
14        self.end_time: float
15        self.elapsed_time: float
16
17    def __enter__(self) -> "TimerContext":
18        self.start_time = time.time()
19        return self
20
21    def __exit__(
22        self,
23        exc_type: type[BaseException] | None,
24        exc_val: BaseException | None,
25        exc_tb: TracebackType | None,
26    ) -> Literal[False]:
27        self.end_time = time.time()
28        self.elapsed_time = self.end_time - self.start_time
29        return False