如何为可以在递归调用中重新分配的 python 函数制定类型提示？

采取以下最小示例：

S = TypeVar("S", bound=int | str)

def meth(a: S) -> S:
    if a == "5":
        return str(meth(int(a)))

    return a

特别是，上面的方法可以采用字符串或整数。它总是返回与其输入相同类型的值，但它可以递归地调用自身，在这种情况下， S 的值可能需要重新分配。使用上面的类型提示，mypy 给出错误：

error: Incompatible return value type (got "str", expected "S")  [return-value]

我如何正确输入上面的函数？

编辑 ： 谢谢，很高兴知道 == 不执行类型缩小。但是，我仅将其用于我的最小示例。我的实际用例的简化版本如下所示：

T = TypeVar("T", bound=CSTNode)

def _rec_set_docstring(expr: T, new_str: str) -> T:
    if isinstance(expr, (BaseSuite, SimpleStatementLine)):
        return expr.with_changes(
            body=(_rec_set_docstring(expr.body[0], new_str),)
        )

    if not isinstance(expr, Expr):
        raise ValueError
    if isinstance(expr.value, (SimpleString, ConcatenatedString)):
        return expr.with_changes(value=expr.value.with_changes(value=new_str))
    raise ValueError

这依赖于 libcs​​t ，但最相关的签名应该是 CSTNode ：

_CSTNodeSelfT = TypeVar("_CSTNodeSelfT", bound="CSTNode")
@dataclass(frozen=True)
class CSTNode(ABC):
    def with_changes(self: _CSTNodeSelfT, **changes: Any) -> _CSTNodeSelfT:
        ...

问题在于你的类型提示告诉 mypy meth 对于同一种类型始终返回相同类型的 S ，但这并非事实。在某些情况下，它返回 str(meth(int(a))) 的结果，该结果始终为 str ，而不管 S 是什么。

你无法在 mypy 当前的类型系统中准确地表达这一点。根本原因是 mypy 不执行类型缩小，因此它不知道在 if 语句的主体中， a 的类型为 str 。

你有几个选择：

1. 忽略错误。 如果确定代码正确，可以安全地忽略类型检查器错误。

```python S = TypeVar("S", bound=int | str)

def meth(a: S) -> S: # type: ignore[return-value] if a == "5": return str(meth(int(a))) return a ```

1. 使用 cast 。 这告诉 mypy 你比它更了解类型。

```python from typing import cast

S = TypeVar("S", bound=int | str)

def meth(a: S) -> S: if a == "5": return cast(S, str(meth(int(a)))) return a ```

1. 重构你的代码。 通常，解决此类类型错误的最佳方法是重构代码，使其更易于进行类型检查。在本例中，你可以将函数拆分为两个单独的函数，一个用于整数，一个用于字符串。

```python def meth_int(a: int) -> int: ...

def meth_str(a: str) -> str: if a == "5": return str(meth_int(int(a))) return a ```

对于你的 _rec_set_docstring 示例，适用相同的原则。mypy 无法推断出 with_changes 在 isinstance 检查后返回的类型与传入的类型相同。你可以通过以下方式解决此问题：

1. 使用 cast ：

```python from typing import cast

T = TypeVar("T", bound=CSTNode)

def _rec_set_docstring(expr: T, new_str: str) -> T: if isinstance(expr, (BaseSuite, SimpleStatementLine)): return cast(T, expr.with_changes( body=(_rec_set_docstring(expr.body[0], new_str),) ))

   if not isinstance(expr, Expr):
       raise ValueError
   if isinstance(expr.value, (SimpleString, ConcatenatedString)):
       return cast(T, expr.with_changes(value=expr.value.with_changes(value=new_str)))
   raise ValueError

```

1. 重构你的代码：

如果没有更具体的例子来说明你的 CSTNode 子类的结构以及你试图实现的目标，很难提供具体的重构建议。但是，一般来说，你可以考虑根据节点类型将函数拆分为更小的、类型更具体的函数。这将使 mypy 更容易跟踪类型并识别潜在问题。