如何消除mypy模块对多类型变量操作时的“不兼容类型”错误？

原问题代码 #

from typing import Union

def f(data: Union[str, bytes]) -> Union[str, bytes]:
    try:
        result = ' '.join(('abc', data))
    except TypeError:
        jointer = ' '.encode('utf-8')
        result = jointer.join(('abc'.encode('utf-8'), data))
    finally:
        return result

mypy报错信息 #

test.py:5: error: Argument 1 to "join" of "str" has incompatible type "Tuple[str, Union[str, bytes]]"; expected "Iterable[str]" [arg-type]
test.py:8: error: Incompatible types in assignment (expression has type "bytes", variable has type "str") [assignment]
test.py:8: error: Argument 1 to "join" of "bytes" has incompatible type "Tuple[bytes, Union[str, bytes]]"; expected "Iterable[Union[bytes, Union[bytearray, memoryview, array[Any], mmap, _CData, PickleBuffer]]]" [arg-type]
Found 3 errors in 1 file (checked 1 source file)

解决方案 #

方法1：使用类型窄化（isinstance检查） #

mypy无法通过try-except推断运行时的类型分支，但通过isinstance明确区分str和bytes类型后，类型检查器可清晰识别各分支的类型：

from typing import Union

def f(data: Union[str, bytes]) -> Union[str, bytes]:
    if isinstance(data, str):
        return ' '.join(('abc', data))
    # mypy会自动推断此处data为bytes类型
    return b' '.join((b'abc', data))

方法2：使用@overload实现精确类型匹配 #

如果需要更精准的输入输出类型关联（比如传入str返回str，传入bytes返回bytes），可以用@overload重载函数定义：

from typing import overload, Union

@overload
def f(data: str) -> str:
    ...

@overload
def f(data: bytes) -> bytes:
    ...

def f(data: Union[str, bytes]) -> Union[str, bytes]:
    if isinstance(data, str):
        return ' '.join(('abc', data))
    return b' '.join((b'abc', data))

错误原因说明 #

1. 原代码try块中str.join要求所有元素为str，但data是Union[str, bytes]，mypy无法保证data一定是str，因此报错。
2. except块中给result赋值bytes，但try块中result被推断为str，类型不兼容。
3. except块的bytes.join要求元素为字节类类型，data的Union类型同样无法满足mypy的类型检查要求。

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