logging.py

"""Logging utilities."""

import contextlib
import json
from datetime import datetime


class LogEntry:
    """A log entry that tracks and formats structured log messages with timing.
    
    Args:
        msg: The main log message
        **kwargs: Additional key-value pairs to include in the log entry
    """
    def __init__(self, msg: str, **kwargs):
        self.start = datetime.now()
        self.entry = {
                         "timestamp": self.start.isoformat(),
                         "message": msg,
                     } | (kwargs or {})
        self.include_execution_time = False

    def __getitem__(self, item):
        return self.entry[item]

    def __setitem__(self, key, value):
        self.entry[key] = value

    def __call__(self, *args, **kwargs):
        self.render(**kwargs)

    def begin(self):
        self["status"] = "begin"
        self.render()
        self.include_execution_time = True
        self["status"] = "interim"

    def complete(self):
        self["status"] = "end"
        self.render()

    def render(self, **kwargs):
        now = datetime.now()
        self.entry["timestamp"] = now.isoformat()
        exec_time = {
            "execution_time": str(now - self.start)
        } if self.include_execution_time else {}
        print(json.dumps(self.entry | (kwargs or {}) | exec_time), flush=True)


@contextlib.contextmanager
def log_it(msg: str, **kwargs):
    """Context manager for logging the execution time of a code block.
    
    Args:
        msg: The message to log
        **kwargs: Additional key-value pairs to include in the log entry
        
    Yields:
        LogEntry: The log entry being used for this context
    """
    entry = LogEntry(msg, **kwargs)
    entry.begin()
    yield entry
    entry.complete()