Source code for llg3d.solvers.profiling

"""Profiling utilities for solvers."""

import time
from functools import wraps
from typing import TYPE_CHECKING, Callable

if TYPE_CHECKING:
    from .base import BaseSolver

ProfilingStats = dict[str, dict[str, float]]




def timeit(func: Callable) -> Callable:
    """
    Decorator to time functions only if profiling is enabled.

    Args:
        func: The function to time.

    Returns:
        A wrapped function that records timing stats to solver.profiling_stats.
    """

    @wraps(func)
    def timeit_wrapper(solver: "BaseSolver", *args, **kwargs):
        # Skip timing overhead if profiling is disabled
        if not solver.profiling:
            return func(solver, *args, **kwargs)

        start = time.perf_counter()
        try:
            return func(solver, *args, **kwargs)
        finally:
            stats = solver.profiling_stats[func.__name__]
            stats["time"] += time.perf_counter() - start
            stats["calls"] += 1

    return timeit_wrapper