muutils.logger.log_util

 1from __future__ import annotations
 2from typing import Any, TypeVar
 3from muutils.jsonlines import jsonl_load_log
 4
 5T_StreamValue = TypeVar("T_StreamValue")
 6
 7
 8def get_any_from_stream(
 9    stream: list[dict[str, T_StreamValue]], key: str
10) -> T_StreamValue:
11    """get the first value of a key from a stream. errors if not found"""
12    for msg in stream:
13        if key in msg:
14            return msg[key]
15
16    raise KeyError(f"key '{key}' not found in stream")
17
18
19def gather_log(file: str) -> dict[str, list[dict[str, Any]]]:
20    """gathers and sorts all streams from a log"""
21    data: list[dict[str, Any]] = jsonl_load_log(file)
22    output: dict[str, list[dict[str, Any]]] = dict()
23
24    for item in data:
25        stream: str = item.get("_stream", "default")
26        if stream not in output:
27            output[stream] = list()
28        output[stream].append(item)
29
30    return output
31
32
33def gather_stream(
34    file: str,
35    stream: str,
36) -> list[dict[str, Any]]:
37    """gets all entries from a specific stream in a log file"""
38    data: list[dict[str, Any]] = jsonl_load_log(file)
39
40    output: list[dict[str, Any]] = list()
41
42    for item in data:
43        # select for the stream
44        if ("_stream" in item) and (item["_stream"] == stream):
45            output.append(item)
46    return output
47
48
49def gather_val(
50    file: str,
51    stream: str,
52    keys: tuple[str, ...],
53    allow_skip: bool = True,
54) -> list[list[Any]]:
55    """gather specific keys from a specific stream in a log file
56
57    example:
58    if "log.jsonl" has contents:
59    ```jsonl
60    {"a": 1, "b": 2, "c": 3, "_stream": "s1"}
61    {"a": 4, "b": 5, "c": 6, "_stream": "s1"}
62    {"a": 7, "b": 8, "c": 9, "_stream": "s2"}
63    ```
64    then `gather_val("log.jsonl", "s1", ("a", "b"))` will return
65    ```python
66    [
67        [1, 2],
68        [4, 5]
69    ]
70    ```
71
72    """
73    data: list[dict[str, Any]] = jsonl_load_log(file)
74
75    output: list[list[Any]] = list()
76
77    for item in data:
78        # select for the stream
79        if ("_stream" in item) and (item["_stream"] == stream):
80            # select for the keys
81            if all(k in item for k in keys):
82                output.append(list(item[k] for k in keys))
83            elif not allow_skip:
84                raise ValueError(f"missing keys '{keys = }' in '{item = }'")
85
86    return output