告诉mypy我肯定知道返回参数的类型

Question

下面的代码从mypy中产生了一个可以理解但错误的错误：

from typing import List, Union, Any
class classA():
def __init__(self, name: str) -> None:
self.__name = name
def __eq__(self, other: Any) -> bool:
if (type(self) == type(other)):
return (self.name == other.name)
return False
@property
def name(self) -> str:
return self.__name

class classB():
def __init__(self, id: int) -> None:
self.__id = id
def __eq__(self, other: Any) -> bool:
if (type(self) == type(other)):
return (self.id == other.id)
return False
@property
def id(self) -> int:
return self.__id

class classC():
def __init__(self) -> None:
self.__elements: List[Union[classA, classB]] = list()
def add(self, elem: Union[classA, classB]) -> None:
if (elem not in self.__elements):
self.__elements.append(elem)
def get_a(self, name_of_a: str) -> classA:
tmp_a = classA(name_of_a)
if (tmp_a in self.__elements):
return self.__elements[self.__elements.index(tmp_a)]
raise Exception(f'{name_of_a} not found')
def get_b(self, id_of_b: int) -> classB:
tmp_b = classB(id_of_b)
if (tmp_b in self.__elements):
return self.__elements[self.__elements.index(tmp_b)]
raise Exception(f'{id_of_b} not found')

调用mypy --show-error-codes classes.py显示以下输出：

classes.py:43: error: Incompatible return value type (got "Union[classA, classB]", expected "classA")  [return-value]
classes.py:49: error: Incompatible return value type (got "Union[classA, classB]", expected "classB")  [return-value]
Found 2 errors in 1 file (checked 1 source file)

如何告诉mypy函数get_a只会返回classA？

Answer 1

您可以使用assert:将此信息告知mypy

def get_a(self, name_of_a: str) -> classA:
tmp_a = classA(name_of_a)
if (tmp_a in self.__elements):
result = self.__elements[self.__elements.index(tmp_a)]
assert isinstance(result, classA)
return result
raise Exception(f'{name_of_a} not found')