如何解决带条件参数函数的mypy类型检查报错问题

问题描述 #

我需要为一个带条件参数的函数添加类型注解：函数要求传入Foo或Bar实例二选一（若foo非None则忽略bar），否则抛出ValueError。但使用mypy检查时出现两类错误：

• 定义foo_or_bar: Union[Foo, Bar]时，提示赋值类型不兼容（表达式类型为Union[Foo, Bar, None]，变量类型为Union[Foo, Bar]）；
• 将类型改为Union[Foo, Bar, None]后，提示Union[Foo, Bar, None]中的Item "None"没有属性some_param_to_check。

原代码示例：

from typing import Optional, Union


class Foo:
    some_param_to_check: str = 'foo_name'
    one_param_exclusive_to_foo: int


class Bar:
    some_param_to_check: str = 'bar_name'
    another_param_exclusive_to_bar: str


def some_process_that_returns_a_bool(
    f_or_b: Union[Foo, Bar],
    a_name: str,
) -> bool:
    return f_or_b.some_param_to_check == a_name


def do_something_with_foo_or_bar(
    foo: Optional[Foo],
    bar: Optional[Bar],
    some_name: str,
) -> bool:
    if not foo and not bar:
        raise ValueError('You need to specify either "foo" or "bar".')

    # 原类型注解引发报错
    foo_or_bar: Union[Foo, Bar]
    foo_or_bar = foo if foo else bar

    return some_process_that_returns_a_bool(foo_or_bar, some_name)


foo_obj = Foo()
bar_obj = Bar()

# 合法调用
do_something_with_foo_or_bar(foo_obj, bar_obj, 'test_string')
do_something_with_foo_or_bar(foo_obj, None, 'test_string')
do_something_with_foo_or_bar(None, bar_obj, 'test_string')

# 非法调用（应抛出异常）
do_something_with_foo_or_bar(None, None, 'test_string')

解决方案 #

方法1：使用assert断言非空 #

因为已经通过if not foo and not bar的判断排除了两者都为None的情况，所以foo_or_bar必然不为None。添加assert可以明确告诉mypy这一点，让它正确推断类型：

def do_something_with_foo_or_bar(
    foo: Optional[Foo],
    bar: Optional[Bar],
    some_name: str,
) -> bool:
    if not foo and not bar:
        raise ValueError('You need to specify either "foo" or "bar".')

    foo_or_bar = foo if foo else bar
    # 断言foo_or_bar非空，让mypy推断其类型为Union[Foo, Bar]
    assert foo_or_bar is not None

    return some_process_that_returns_a_bool(foo_or_bar, some_name)

方法2：重构分支逻辑，让mypy自动推断类型 #

改用显式的分支判断替代三元表达式，mypy会根据分支直接推断foo_or_bar的类型为Union[Foo, Bar]：

def do_something_with_foo_or_bar(
    foo: Optional[Foo],
    bar: Optional[Bar],
    some_name: str,
) -> bool:
    if not foo and not bar:
        raise ValueError('You need to specify either "foo" or "bar".')

    # 显式分支让mypy直接推断类型
    if foo:
        foo_or_bar = foo
    else:
        # 前面已排除两者都为None的情况，bar必然非空
        foo_or_bar = bar

    return some_process_that_returns_a_bool(foo_or_bar, some_name)

方法3：使用typing.cast强制类型转换 #

如果不想修改逻辑，可以用cast明确告诉mypyfoo_or_bar的实际类型：

from typing import Optional, Union, cast

# ... 其他代码不变 ...

def do_something_with_foo_or_bar(
    foo: Optional[Foo],
    bar: Optional[Bar],
    some_name: str,
) -> bool:
    if not foo and not bar:
        raise ValueError('You need to specify either "foo" or "bar".')

    # 用cast强制转换类型为Union[Foo, Bar]
    foo_or_bar = cast(Union[Foo, Bar], foo if foo else bar)

    return some_process_that_returns_a_bool(foo_or_bar, some_name)

说明 #

三种方法都能让mypy通过类型检查，同时保留代码的原有逻辑。其中方法2的分支逻辑更清晰，mypy的推断也更自然，是推荐的方案。

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