如何让PyCharm识别类装饰器添加的类变量logger？

问题描述 #

我编写了一个给类自动添加日志功能的类装饰器，代码运行正常，但PyCharm、mypy等类型检查工具都会报错，提示被装饰的类不存在logger属性，触发警告的代码行是访问Spam.logger.name的部分。

最初的代码如下：

from typing import TypeVar

_T = TypeVar("_T", bound=type)


def with_logger(cls: _T) -> _T:
    from logging import getLogger

    cls.logger = getLogger(cls.__name__)
    return cls


@with_logger
class Spam:
    def __init__(self) -> None:
        print(f"{Spam.logger.name = }")


Spam()

尝试过的无效方案 #

1. 通过Protocol和动态类修改装饰器：参考相关问题修改后，类型检查工具仍无法识别logger属性：

from typing import TypeVar

_T = TypeVar("_T", bound=type)


def with_logger(cls: _T) -> _T:
    from logging import getLogger
    from typing import TYPE_CHECKING

    setattr(cls, 'logger', getLogger(cls.__name__))
    if TYPE_CHECKING:
        from typing import ClassVar, Protocol, cast

        class LoggerMixin(Protocol):
            from logging import Logger

            logger: ClassVar[Logger] = getLogger(cls.__name__)

        return cast(_T, type(cls.__name__, (LoggerMixin, cls), dict()))
    else:
        return cls


@with_logger
class Spam:
    def __init__(self) -> None:
        print(f"{Spam.logger.name = }")


Spam()

2. 在被装饰类显式标注类型：此方法能解决警告，但需要每个类手动添加标注，失去了装饰器简化代码的作用：

from logging import Logger, getLogger
from typing import ClassVar, TypeVar

_T = TypeVar("_T", bound=type)


def with_logger(cls: _T) -> _T:
    setattr(cls, 'logger', getLogger(cls.__name__))
    return cls


@with_logger
class Spam:
    logger: ClassVar[Logger]

    def __init__(self) -> None:
        print(f"{Spam.logger.name = }")


Spam()

3. 继承基类方案：该方案可行且类型检查工具能识别，但希望用装饰器实现：

from logging import Logger, getLogger
from typing import ClassVar


class BaseLogger:
    logger: ClassVar[Logger]
    
    def __new__(cls):
        cls.logger = getLogger(cls.__name__)
        return super().__new__(cls)


class Spam(BaseLogger):
    def __init__(self) -> None:
        print(f"{Spam.logger.name = }")


Spam()

可行的装饰器解决方案 #

通过定义包含logger属性的Protocol，并让装饰器返回原类类型与该Protocol的交集类型，即可让类型检查工具识别装饰器添加的logger属性：

from logging import Logger
from typing import Protocol, ClassVar, TypeVar, cast

# 定义包含logger的协议
class HasLogger(Protocol):
    logger: ClassVar[Logger]

_T = TypeVar("_T", bound=type)

def with_logger(cls: _T) -> _T & HasLogger:
    from logging import getLogger
    # 给类添加logger属性
    cls.logger = getLogger(cls.__name__)
    # 类型转换，告诉类型检查工具该类现在具备HasLogger的属性
    return cast(_T & HasLogger, cls)

@with_logger
class Spam:
    def __init__(self) -> None:
        print(f"{Spam.logger.name = }")

Spam()

说明 #

• HasLogger协议定义了类必须包含ClassVar[Logger]类型的logger属性，类型检查工具会认可实现该协议的类具备此属性。
• 装饰器的返回类型_T & HasLogger表示返回的类既是原类型_T，同时也符合HasLogger协议，这样类型检查工具就能识别logger属性。
• 此方案无需修改被装饰的类，完全通过装饰器实现类型提示，同时保持代码简洁。

内容的提问来源于stack exchange，提问作者StSav012