如何让自定义断言函数被Python类型检查器识别？

要让my_assert这类自定义断言函数被PyLance/Pyright等静态类型检查器识别，实现和内置assert一样的类型窄化效果，可通过以下轻量方式解决：

方法1：标注函数返回类型为Literal[True] #

将自定义断言函数的返回类型标注为Literal[True]，告知类型检查器：只有当传入的condition为True时，函数才会正常返回。调用函数后，检查器会自动根据condition的表达式窄化变量类型。

from typing import Literal
import logging

logger = logging.getLogger(__name__)

def my_assert(condition: bool, msg: str) -> Literal[True]:
    if not condition:
        logger.error(msg)
        raise AssertionError(msg)
    return True

# 使用示例
def add_one(a: int | None) -> int:
    my_assert(a is not None, "a is None")
    return a + 1  # Pylance不再提示a可能为None

方法2：结合NoReturn标注异常分支返回类型 #

标注函数在条件不满足时会抛出异常且不会返回（NoReturn），让类型检查器明确：执行到函数调用后的代码时，condition必然为True。

from typing import NoReturn
import logging

logger = logging.getLogger(__name__)

def my_assert(condition: bool, msg: str) -> None | NoReturn:
    if not condition:
        logger.error(msg)
        raise AssertionError(msg)

# 使用示例
def add_one(a: int | None) -> int:
    my_assert(a is not None, "a is None")
    return a + 1  # 类型检查器自动识别a为int类型

方法3：利用类型守卫实现类型专属断言 #

如果需要针对特定类型做断言，可使用TypeGuard创建类型专属的断言函数，类型检查精度更高，适合有明确类型约束的场景。

from typing import TypeVar, TypeGuard
import logging

logger = logging.getLogger(__name__)
T = TypeVar('T')

def assert_not_none(obj: T | None, msg: str) -> TypeGuard[T]:
    if obj is None:
        logger.error(msg)
        raise AssertionError(msg)
    return True

# 使用示例
def add_one(a: int | None) -> int:
    assert_not_none(a, "a is None")
    return a + 1  # 类型检查器直接识别a为int

以上方式仅需调整自定义断言函数的类型注解，无需大幅修改原有业务代码，即可让静态类型检查器正确识别断言逻辑，消除不必要的类型提示错误。

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