import functools
import time


def timeit(runtime_attr: str):
    """
    Decorate a function to compute its run time and update the \
    :class:`IRLSState` instance attribute called ``runtime_attr`` with the run \
    time of the function.

    The first argument of ``fct`` should be an :class:`IRLSState` instance.
    """

    def fct_wrap(fct):
        @functools.wraps(fct)
        def inner_fct(*args, **kwargs):
            start = time.perf_counter()
            out = fct(*args, **kwargs)
            setattr(args[0], runtime_attr, time.perf_counter() - start)
            return out

        return inner_fct

    return fct_wrap