How to change logger depth to go "down" one level?

[Jerry-Ma]

Hi,
I would like to create a helper function that can record the duration of a function call, like this:

from loguru import logger as _loguru_logger
from contextlib import ContextDecorator
import time


class timeit(ContextDecorator):

    """Context decorator that logs the execution time of the decorated item.

    Parameters
    ----------
    arg: object or str
        If `arg` type is `str`, a new decorator is returned which uses `arg`
        in the generated message in places of the name of the decorated object.
    """

    _logger = _loguru_logger.patch(
        lambda r: r.update(name=f'timeit: {r["name"]}'),
    )

    def __new__(cls, arg=None, **kwargs):
        if callable(arg):
            return cls(arg.__name__, **kwargs)(arg)
        return super().__new__(cls)

    def __init__(self, msg=None, level="DEBUG"):
        self._msg = msg
        self._level = level
        self._logger = self._logger.opt(depth=1)

    def __enter__(self):
        self._logger.log(self._level, f"{self._msg} ...")
        self._start = time.time()

    def __exit__(self, *args):
        elapsed = time.time() - self._start
        self._logger.log(
            self._level,
            f"{self._msg} done in {self._format_time(elapsed)}",
        )

    def __call__(self, arg, *args, **kwargs):
        self._logger = self._logger.opt(depth=2)
        if callable(arg):
            self._msg = arg.__name__
        return super().__call__(arg, *args, **kwargs)

    @staticmethod
    def _format_time(time):
        max_ms = 15
        if time < max_ms:
            return f"{time * 1e3:.0f}ms"
        return f"{time:.2f}s"


@timeit
def some_func1():
    pass


@timeit(level="INFO")
def some_func2():
    pass


def outer():
    some_func1()


if __name__ == "__main__":

    some_func1()

    some_func2()

    with timeit("some_block"):
        pass

    outer()

Right now, the output looks like this:

2024-04-05 20:43:54.843 | DEBUG    | timeit: __main__:<module>:73 - some_func1 ...
2024-04-05 20:43:54.843 | DEBUG    | timeit: __main__:<module>:73 - some_func1 done in 0ms
2024-04-05 20:43:54.843 | INFO     | timeit: __main__:<module>:75 - some_func2 ...
2024-04-05 20:43:54.843 | INFO     | timeit: __main__:<module>:75 - some_func2 done in 0ms
2024-04-05 20:43:54.843 | DEBUG    | timeit: __main__:<module>:77 - some_block ...
2024-04-05 20:43:54.843 | DEBUG    | timeit: __main__:<module>:77 - some_block done in 0ms
2024-04-05 20:43:54.843 | DEBUG    | timeit: __main__:outer:68 - some_func1 ...
2024-04-05 20:43:54.843 | DEBUG    | timeit: __main__:outer:68 - some_func1 done in 0ms

However, for the cases that I use the timeit function as decorator, I would like the module name to be bind to the function itself, not the context that defines the function, that is, make the code behaves like the following:

def some_func1()
    logger.debug("start")
    logger.debug("stop")

This would mean to go "down" one level in the frames, but I am not sure how to do that. Setting opt(depth=-1) does not help because it will be selelcting the bottom of the caller stack.

Any ideas?

[Delgan]

That's not really feasible, because the logger is being used before the frame for the wrapped function even exist. They are not part of the same call stack.

I'd recommend leaving it as it is. Otherwise, it would be confusing and inaccurate anyway: the logging call is made within timeit(), not in func1().

If it really bothers you, I think you can get away with using a custom format and bind(), or by changing the function value via logger.patch().