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

Question

如果可以建立连接，我有一个返回 psycopg2 连接的函数。所以返回类型应该是

Optional[psycopg2.connection]

或

psycopg2.connection | None

。但是我无法在运行时导入

psycopg2.connection

。我已经尝试了如何从类型化存根文件导入类型定义？中提到的解决方法，但这给了我这个 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

Answer 1

您链接的问题提出了一些极其肮脏的黑客，但它似乎不再起作用。这么简单的情况下根本就没有必要。此外，说实话，我无法在从

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(...)

未来导入的文档。这与直接使用字符串非常相似，但看起来更好并且更方便，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

，如上所述以保持一致性。但是，

annotations

future-import 使这个表达式有效地成为一个字符串，因此如果您的工具不执行任何运行时类型自省 - 您仍然可以在旧版本上使用管道联合语法。请注意，在 3.10 之前的 Python 上，使用此语法

typing.get_type_hints

和类似实用程序将会失败。

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