使用 psycopg2 类型存根注释返回类型

如果可以建立连接，我有一个返回 psycopg2 连接的函数。所以返回类型应该是Optional[psycopg2.connection], or psycopg2.connection | None。但是我无法导入psycopg2.connection在运行时。我已经尝试过中提到的解决方法如何从 typeshed 存根文件导入类型定义？ https://stackoverflow.com/questions/66007912/how-can-i-import-type-definitions-from-a-typeshed-stub-file但这给了我这个 mypy 错误：Single overload definition, multiple required。这是我的代码

import psycopg2
from typing import Optional, TYPE_CHECKING, overload

if TYPE_CHECKING:
    from psycopg2 import connection
    
    @overload
    def get_connection() -> Optional[connection]: ...


# Make DB error logging less spammy
has_logged_error = False

def get_connection():
    try:
        conn = psycopg2.connect(
            dbname=settings.db_name,
            user=settings.db_user,
            password=settings.db_password,
            host=settings.db_host,
            port=settings.db_port,
        )
        return conn
    except Exception as e:
        global has_logged_error
        if not has_logged_error:
            logger.error(f"Error connecting to DB: {e}")
            has_logged_error = True
        return

您链接的问题提出了一些极其肮脏的黑客，但它似乎不再起作用。这么简单的情况下根本就没有必要。此外，说实话，我无法在任何mypy版本从0.800（足够老了，考虑到链接的答案是最近的），所以也许从来没有工作过。

为了便于阅读，我减少了您的代码示例，使其仅包含最小的返回值。

变体 1：使用条件导入和字符串注释

import psycopg2
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from psycopg2 import connection
    
def get_connection() -> Optional['connection']:
    return psycopg2.connect(...)

这很简单：mypy知道什么connection是（在存根中定义）；运行时不会尝试学习一些东西connection因为它只看到一个字符串。

变体 2：使用条件导入和注释 future

from __future__ import annotations
import psycopg2
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from psycopg2 import connection
    
def get_connection() -> Optional[connection]:
    return psycopg2.connect(...)

未来进口的文件 https://docs.python.org/3/library/__future__.html。这与直接使用字符串非常相似，但看起来更好并且更方便，IMO。

变体 3：使用字符串注释，但避免条件导入

import psycopg2
from typing import Optional
    
def get_connection() -> Optional['psycopg2.connection']:
    return psycopg2.connect(...)

变体 4：使用 future 注释，但避免条件导入

from __future__ import annotations
import psycopg2
from typing import Optional
    
def get_connection() -> Optional[psycopg2.connection]:
    return psycopg2.connect(...)

变体 3 和 4 没有暴露这一点connection仅存根，将其隐藏为实现细节。您可能更愿意明确说明这一点 - 然后使用 1 或 2。

修改以使用当前功能

这是我最喜欢的。 Union 语法在 python 3.10+ 中有效，因此如果您使用较旧的语法 - 您可能需要坚持使用Optional如上所述，以保持一致性。然而，annotationsfuture-import 使这个表达式有效地成为一个字符串，因此如果您的工具不执行任何运行时类型自省 - 您仍然可以在旧版本上使用管道联合语法。请注意typing.get_type_hints在 3.10 之前的 python 上，类似的实用程序将因此语法而失败。

from __future__ import annotations
import psycopg2
    
def get_connection() -> psycopg2.connection | None:
    return psycopg2.connect(...)