Python TypeGuard类型收窄：TypeVar返回类型匹配报错解决

场景示例 #

from typing import TypeVar
from typing_extensions import TypeGuard

_T = TypeVar('_T', str, int)

def is_str(val: _T) -> TypeGuard[str]:
    return isinstance(val, str)

def is_int(val: _T) -> TypeGuard[int]:
    return isinstance(val, int)

def process_str_int(data: _T) -> _T:
    if is_str(data):
        # 原注释有误，实际此处data被收窄为str
        print("Returning a string")
        return data
    elif is_int(data):
        print("Returning an int")
        return data

def process_str_int_with_isinstance(data: _T) -> _T:
    if isinstance(data, str):
        # 原注释有误，实际此处data被收窄为str
        print("Returning a string")
        return data
    elif isinstance(data, int):
        print("Returning an int")
        return data

process_str_int("hello")

问题描述 #

在Python 3.9环境下，process_str_int函数的return语句处，Pyright报错：'Expression of type "str" cannot be assigned to return type "_T@process_str_int"'（int分支同理）。而使用isinstance的process_str_int_with_isinstance函数无此问题，TypeVar可被正确收窄，返回类型能被正确推断并匹配返回值。需要实现自定义TypeGuard的正确类型收窄行为。

解决方案 #

问题根源在于原TypeGuard仅断言值为某一类型，但未关联到受约束的TypeVar。修改自定义TypeGuard的返回类型，明确断言值同时属于TypeVar和目标类型，即可解决类型不匹配问题。

修改后的代码 #

from typing import TypeVar
from typing_extensions import TypeGuard

_T = TypeVar('_T', str, int)

def is_str(val: _T) -> TypeGuard[_T & str]:
    return isinstance(val, str)

def is_int(val: _T) -> TypeGuard[_T & int]:
    return isinstance(val, int)

def process_str_int(data: _T) -> _T:
    if is_str(data):
        print("Returning a string")
        return data  # 类型收窄为_T & str，匹配返回类型_T
    elif is_int(data):
        print("Returning an int")
        return data  # 类型收窄为_T & int，匹配返回类型_T

def process_str_int_with_isinstance(data: _T) -> _T:
    if isinstance(data, str):
        print("Returning a string")
        return data
    elif isinstance(data, int):
        print("Returning an int")
        return data

process_str_int("hello")

原理说明 #

原写法TypeGuard[str]仅告知类型检查器val是str，但无法确认该str属于_T约束的分支，因此认为返回值与_T类型不兼容。而TypeGuard[_T & str]明确了val同时是_T和str的实例，既完成了类型收窄，又保留了与原TypeVar的关联，让类型检查器认可返回值符合_T的要求。

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