如何对 Python 函数进行类型注释，以仅在输入为 None 时将其标记为返回 None？

Question

我们的代码中有一个典型的数据转换模式：当值为 None 时，我们让它通过。例如，

def capitalize(value):  
    if value is None:  
        return None  
    return value.capitalize()
     
# usage example:
assert capitalize(None) is None
assert capitalize("hello world") == "Hello world"

我可以这样注释它：

from typing import Optional


def capitalize(value: Optional[str]) -> Optional[str]:  
    if value is None:  
        return None  
    return value.capitalize()

看起来不错，但是下面的代码

capitalize("Hello world").split()

总是会让mypy抱怨。

error: Item "None" of "Optional[str]" has no attribute "split"

有没有办法用类型注解来表达“None总是转换为None，str总是转换为str”的转换规则？

Answer 1

这听起来像是具有值限制的泛型的用例。

下面的代码基本上说T可以是 astr或None，函数定义说“这个函数返回与传入的类型相同的类型”。

from typing import TypeVar

T = TypeVar("T", str, None)

def capitalize(value: T) -> T:
    if value is None:
        return None
    return value.capitalize()

capitalize("Hello world").split()

在上面的代码上运行mypy似乎工作正常，并且：

capitalize(None).split()

抱怨的原因mypy：error: "None" has no attribute "split"我认为这就是你所追求的。