utils/timing.py

import inspect
import time
from datetime import datetime
from functools import wraps
from typing import Any, Callable

import streamlit as st


def summarize_parameter_value(value: str, max_length: int = 100) -> str:
    """Summarize parameter values that are too long or complex."""
    if not value:
        return ""

    # Handle DataFrames
    if "DataFrame" in value and "[" in value and "]" in value:
        try:
            dims = value[value.find("[") + 1 : value.find("]")]
            return f"DataFrame[{dims}]"
        except Exception:
            return "DataFrame"

    # Handle lists, tuples, and other sequences
    if value.startswith(("[", "(", "{")):
        try:
            item_count = value.count(",") + 1
            return f"{value[:20]}... ({item_count} items)"
        except Exception:
            return f"{value[:20]}..."

    # Handle long strings
    if len(value) > max_length:
        return f"{value[:max_length]}..."

    return value


def timer(include_params: bool = False) -> Callable:
    """
    Decorator to time function execution and store results in session state.

    Args:
        include_params: Whether to include function parameters in timing logs
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(*args: Any, **kwargs: Any) -> Any:
            if not st.session_state.get("ENABLE_TIMING", False):
                return func(*args, **kwargs)

            start = time.perf_counter()
            result = func(*args, **kwargs)
            end = time.perf_counter()
            duration = end - start

            # Initialize timing stats if needed
            if "timing_stats" not in st.session_state:
                st.session_state.timing_stats = {}
                st.session_state.timing_logs = []

            # Initialize list for this function if needed
            if func.__name__ not in st.session_state.timing_stats:
                st.session_state.timing_stats[func.__name__] = []

            # Append new duration
            st.session_state.timing_stats[func.__name__].append(duration)

            # Create log entry
            log_entry = {
                "timestamp": datetime.now().isoformat(),
                "function": func.__name__,
                "duration": duration,
            }

            if include_params:
                # Get parameter names from function signature
                sig = inspect.signature(func)
                param_names = list(sig.parameters.keys())

                # Combine args and kwargs into a parameter dictionary
                param_values = {}
                for i, arg in enumerate(args):
                    if i < len(param_names):
                        param_values[param_names[i]] = summarize_parameter_value(
                            str(arg),
                            max_length=40,
                        )
                param_values.update(
                    {
                        k: summarize_parameter_value(str(v), max_length=40)
                        for k, v in kwargs.items()
                    }
                )

                log_entry["parameters"] = param_values

            st.session_state.timing_logs.append(log_entry)
            return result

        return wrapper

    return decorator