正确输入带有可选关键字参数的装饰器

问题描述 投票:0回答:0

我想在 Python 3.11 中编写一个装饰器,为函数添加一些基本的日志记录。

我想使用没有任何关键字参数的装饰器:

@add_logging

在那种情况下,它应该使用默认的日志记录级别

logging.DEBUG
.

我还想使用带有关键字参数的装饰器,然后指定日志记录级别:

@add_logging(logging_level=logging.ERROR)

所有这些都应该最大限度地打字。

我想出了如何使用 @overload 装饰器来做到这一点:

P = ParamSpec("P")
R = TypeVar("R")


@overload
def add_logging(function: Callable[P, R]) -> Callable[P, R]:
    ...


@overload
def add_logging(*, logging_level: int = DEBUG) -> Callable[[Callable[P, R]], Callable[P, R]]:
    ...


def add_logging(
    function: Optional[Callable[P, R]] = None, *, logging_level: int = DEBUG
) -> Callable[[Callable[P, R]], Callable[P, R]] | Callable[P, R]:
    """A type-safe decorator to add logging to a function.

    """

    def wrapper(wrapped_function: Callable[P, R], *args: P.args, **kwargs: P.kwargs) -> R:
        module_name: str = wrapped_function.__module__
        logger: Logger = getLogger(module_name)

        if not logger.isEnabledFor(logging_level):
            return wrapped_function(*args, **kwargs)

        filename: str = wrapped_function.__code__.co_filename
        first_line_no: int = wrapped_function.__code__.co_firstlineno

        logger.handle(
            LogRecord(
                name=module_name,
                level=logging_level,
                pathname=filename,
                lineno=first_line_no,
                msg="Call *%r **%r",
                args=(args, kwargs),
                exc_info=None,
                func=wrapped_function.__qualname__,
            )
        )

        try:
            result = wrapped_function(*args, **kwargs)
        except Exception as exception:
            logger.handle(
                LogRecord(
                    name=module_name,
                    level=logging_level,
                    pathname=filename,
                    lineno=first_line_no,
                    msg="",
                    args=None,
                    exc_info=(type(exception), exception, None),
                    func=wrapped_function.__qualname__,
                )
            )
            raise

        logger.handle(
            LogRecord(
                name=module_name,
                level=logging_level,
                pathname=filename,
                lineno=first_line_no,
                msg="Return %r",
                args=(result,),
                exc_info=None,
                func=wrapped_function.__qualname__,
            )
        )
        return result

    # Without arguments, `function` is passed directly to the decorator
    if function is not None:
        if not callable(function):
            raise TypeError(f"Expected positional parameter of type callable, but found type {type(function)} instead.")
        return wraps(function)(partial(wrapper, function))

    # With arguments, we need to return a function that accepts the function
    def decorator(function_with_args: Callable[P, R]) -> Callable[P, R]:
        return wraps(function_with_args)(partial(wrapper, function_with_args))

    return decorator

这被 mypy 完全接受并按上述方式工作,即有和没有关键字参数

logging_level
。这是一个用这个装饰的函数的示例输出:

[2023-05-01 00:08:19 +0100] [__main__] [add_two_ints] [DEBUG] Call *(5, 7) **{}
[2023-05-01 00:08:19 +0100] [__main__] [add_two_ints] [DEBUG] Return 12

我使用 LogRecord 而不是 logging.debug/info/etc 的原因是日志记录级别的灵活性以及显式设置包装函数名称的能力,因为这是我的日志记录策略的一部分。

然而,代码看起来相当臃肿。我在这里错过了什么吗?有没有更简单、更优雅、更明显的方法来实现这一目标?

请注意,我正在寻找完整类型的代码示例,因为这是我在尝试解决此问题时面临的挑战之一。

python-3.x python-decorators mypy
© www.soinside.com 2019 - 2024. All rights reserved.