为返回psycopg2连接的函数标注返回类型遇mypy报错求解决

我有一个函数，成功建立连接时会返回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

错误原因 #

@overload装饰器用于标记函数存在多个不同签名（比如不同参数组合或返回类型），但你的函数只有一种返回逻辑，仅定义一个重载会触发mypy的Single overload definition, multiple required错误。

修正方案 #

移除不必要的@overload，直接利用TYPE_CHECKING的条件导入添加类型提示，同时明确返回None让代码更规范：

方案1（兼容Python 3.9及以下）

import psycopg2
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from psycopg2 import connection

has_logged_error = False

def get_connection() -> Optional["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 None

方案2（Python 3.10+ 支持联合类型）

import psycopg2
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from psycopg2 import connection

has_logged_error = False

def get_connection() -> "connection | None":
    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 None

关键说明 #

• 在TYPE_CHECKING块中导入connection类型，仅会在类型检查阶段生效，不会影响运行时
• 用字符串引号包裹类型（如"connection"），避免运行时解析类型导致的导入问题
• 将return改为return None，让代码逻辑更清晰，同时严格符合返回类型的定义

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