如何注释返回类型取决于其参数的函数？

Posted 2023-02-19

技术标签:

【中文标题】如何注释返回类型取决于其参数的函数？【英文标题】：How do I annotate a function whose return type depends on its argument? 【发布时间】：2021-10-15 01:38:49 【问题描述】：

在 Python 中，我经常编写过滤集合以查找特定子类型实例的函数。例如，我可能会在 DOM 中查找特定类型的节点或在日志中查找特定类型的事件：

def find_pre(soup: TagSoup) -> List[tags.pre]:
    """Find all <pre> nodes in `tag_soup`."""
    …

def filter_errors(log: List[LogEvent]) -> List[LogError]:
    """Keep only errors from `log`.""" 
    …

为这些函数编写类型很容易。但是这些函数的泛型版本需要一个参数来指定要返回的类型呢？

def find_tags(tag_soup: TagSoup, T: type) -> List[T]:
    """Find all nodes of type `T` in `tag_soup`."""
    …

def filter_errors(log: List[LogEvent], T: type) -> List[T]:
    """Keep only events of type `T` from `log`.""" 
    …

（上面的签名是错误的：我不能在返回类型中引用T。）

这是一个相当常见的设计：docutils 有node.traverse(T: type)，BeautifulSoup 有soup.find_all() 等。当然它可以变得任意复杂，但是Python 类型注释可以处理像上面这样的简单情况吗？

这里有一个 MWE 让它变得非常具体：

from dataclasses import dataclass
from typing import *

@dataclass
class Packet: pass

@dataclass
class Done(Packet): pass

@dataclass
class Exn(Packet):
    exn: str
    loc: Tuple[int, int]

@dataclass
class Message(Packet):
    ref: int
    msg: str

Stream = Callable[[], Union[Packet, None]]

def stream_response(stream: Stream, types) -> Iterator[??]:
    while response := stream():
        if isinstance(response, Done): return
        if isinstance(response, types): yield response

def print_messages(stream: Stream):
    for m in stream_response(stream, Message):
        print(m.msg) # Error: Cannot access member "msg" for "Packet"

msgs = iter((Message(0, "hello"), Exn("Oops", (1, 42)), Done()))
print_messages(lambda: next(msgs))

Pyright 说：

  29:17 - error: Cannot access member "msg" for type "Packet"
  Member "msg" is unknown (reportGeneralTypeIssues)

在上面的例子中，有没有办法注解stream_response，以便Python类型检查器接受print_messages的定义？

【问题讨论】：

打字文档的this section 有帮助吗？ TypeVar 似乎正是您所需要的。 @Kemp：它没有：在def f(x: T) -> List[T] 中，返回类型取决于x 的type。在def f(x: type) -> List[x]（我想要/需要的）中，返回类型取决于x 的值。 @AlexWaygood：不，返回类型更精确：首先，它永远不是None；其次，它保证是特定类型的数据包（以types中传递的为准。 @AlexWaygood 后者最好，但作为最后的手段，前者也可以。实际上，在我的情况下不会超过 10 个；我很想知道它会扮演什么角色！ 【参考方案1】：

好的，我们开始吧。它通过了 MyPy --strict，但它并不漂亮。

这里发生了什么

对于给定的类A，我们知道A 的实例类型将是A（显然）。但是A 本身的类型是什么？从技术上讲，A 的类型是type，因为所有不使用元类的python 类都是type 的实例。然而，用type 注释一个参数并不能告诉类型检查器太多。相反，用于在类型层次结构中“上一层”的 Python 类型检查语法是 Type[A]。因此，如果我们有一个函数myfunc，它返回一个作为参数输入的类的实例，我们可以相当简单地注释如下：

from typing import TypeVar, Type

T = TypeVar('T')

def myfunc(some_class: Type[T]) -> T:
    # do some stuff
    return some_class()

但是，您的情况要复杂得多。您可以输入一个类作为参数，或者您可以输入两个类或三个类......等等。我们可以使用typing.overload 解决这个问题，它允许我们为给定函数注册多个签名。这些签名在运行时被完全忽略；它们纯粹用于类型检查器；因此，这些函数的主体可以留空。通常，您只在用@overload 装饰的函数体中放置一个文档字符串或文字省略号...。

我认为没有办法概括这些重载函数，这就是为什么可以传递给types 参数的最大元素数很重要。您必须繁琐地枚举函数的每个可能签名。如果您沿着这条路线走，您可能需要考虑将 @overload 签名移动到单独的 .pyi 存根文件中。

from dataclasses import dataclass
from typing import (
    Callable,
    Tuple,
    Union,
    Iterator,
    overload,
    TypeVar,
    Type, 
    Sequence
)

@dataclass
class Packet: pass

P1 = TypeVar('P1', bound=Packet)
P2 = TypeVar('P2', bound=Packet)
P3 = TypeVar('P3', bound=Packet)
P4 = TypeVar('P4', bound=Packet)
P5 = TypeVar('P5', bound=Packet)
P6 = TypeVar('P6', bound=Packet)
P7 = TypeVar('P7', bound=Packet)
P8 = TypeVar('P8', bound=Packet)
P9 = TypeVar('P9', bound=Packet)
P10 = TypeVar('P10', bound=Packet)

@dataclass
class Done(Packet): pass

@dataclass
class Exn(Packet):
    exn: str
    loc: Tuple[int, int]

@dataclass
class Message(Packet):
    ref: int
    msg: str

Stream = Callable[[], Union[Packet, None]]

@overload
def stream_response(stream: Stream, types: Type[P1]) -> Iterator[P1]:
    """Signature if exactly one type is passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2]]
) -> Iterator[Union[P1, P2]]:
    """Signature if exactly two types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2], Type[P3]]
) -> Iterator[Union[P1, P2, P3]]:
    """Signature if exactly three types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2], Type[P3], Type[P4]]
) -> Iterator[Union[P1, P2, P3, P4]]:
    """Signature if exactly four types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2], Type[P3], Type[P4], Type[P5]]
) -> Iterator[Union[P1, P2, P3, P4, P5]]:
    """Signature if exactly five types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2], Type[P3], Type[P4], Type[P5], Type[P6]]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6]]:
    """Signature if exactly six types are passed in for the `types` parameter"""

@overload
def stream_response(
    stream: Stream, 
    types: Tuple[
        Type[P1], 
        Type[P2],
        Type[P3],
        Type[P4], 
        Type[P5],
        Type[P6],
        Type[P7]
    ]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6, P7]]:
    """Signature if exactly seven types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[
        Type[P1], 
        Type[P2],
        Type[P3],
        Type[P4], 
        Type[P5],
        Type[P6],
        Type[P7],
        Type[P8]
    ]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6, P7, P8]]:
    """Signature if exactly eight types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[
        Type[P1], 
        Type[P2],
        Type[P3],
        Type[P4], 
        Type[P5],
        Type[P6],
        Type[P7],
        Type[P8],
        Type[P9]
    ]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6, P7, P8, P9]]:
    """Signature if exactly nine types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[
        Type[P1], 
        Type[P2],
        Type[P3],
        Type[P4], 
        Type[P5],
        Type[P6],
        Type[P7],
        Type[P8],
        Type[P9],
        Type[P10]
    ]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6, P7, P8, P9, P10]]:
    """Signature if exactly ten types are passed in for the `types` parameter"""

# We have to be more generic in our type-hinting for the concrete implementation 
# Otherwise, MyPy struggles to figure out that it's a valid argument to `isinstance`
def stream_response(
    stream: Stream,
    types: Union[type, Tuple[type, ...]]
) -> Iterator[Packet]:
    
    while response := stream():
        if isinstance(response, Done): return
        if isinstance(response, types): yield response

def print_messages(stream: Stream) -> None:
    for m in stream_response(stream, Message):
        print(m.msg)

msgs = iter((Message(0, "hello"), Exn("Oops", (1, 42)), Done()))
print_messages(lambda: next(msgs))

减少冗长的策略

如果您想让它更简洁，实现此目的的一种方法是为某些类型结构引入别名。这里的危险是类型提示的意图和含义变得很难阅读，但它确实使重载 7-10 看起来不那么可怕：

from dataclasses import dataclass
from typing import (
    Callable,
    Tuple,
    Union,
    Iterator,
    overload,
    TypeVar,
    Type, 
    Sequence
)

@dataclass
class Packet: pass

P1 = TypeVar('P1', bound=Packet)
P2 = TypeVar('P2', bound=Packet)
P3 = TypeVar('P3', bound=Packet)
P4 = TypeVar('P4', bound=Packet)
P5 = TypeVar('P5', bound=Packet)
P6 = TypeVar('P6', bound=Packet)
P7 = TypeVar('P7', bound=Packet)
P8 = TypeVar('P8', bound=Packet)
P9 = TypeVar('P9', bound=Packet)
P10 = TypeVar('P10', bound=Packet)

_P = TypeVar('_P', bound=Packet)
S = Type[_P]

T7 = Tuple[S[P1], S[P2], S[P3], S[P4], S[P5], S[P6], S[P7]]
T8 = Tuple[S[P1], S[P2], S[P3], S[P4], S[P5], S[P6], S[P7], S[P8]]
T9 = Tuple[S[P1], S[P2], S[P3], S[P4], S[P5], S[P6], S[P7], S[P8], S[P9]]
T10 = Tuple[S[P1], S[P2], S[P3], S[P4], S[P5], S[P6], S[P7], S[P8], S[P9], S[P10]]

@dataclass
class Done(Packet): pass

@dataclass
class Exn(Packet):
    exn: str
    loc: Tuple[int, int]

@dataclass
class Message(Packet):
    ref: int
    msg: str

Stream = Callable[[], Union[Packet, None]]

@overload
def stream_response(stream: Stream, types: Type[P1]) -> Iterator[P1]:
    """Signature if exactly one type is passed in for the `types` parameter"""

@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2]]
) -> Iterator[Union[P1, P2]]:
    """Signature if exactly two types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2], Type[P3]]
) -> Iterator[Union[P1, P2, P3]]:
    """Signature if exactly three types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2], Type[P3], Type[P4]]
) -> Iterator[Union[P1, P2, P3, P4]]:
    """Signature if exactly four types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2], Type[P3], Type[P4], Type[P5]]
) -> Iterator[Union[P1, P2, P3, P4, P5]]:
    """Signature if exactly five types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: Tuple[Type[P1], Type[P2], Type[P3], Type[P4], Type[P5], Type[P6]]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6]]:
    """Signature if exactly six types are passed in for the `types` parameter"""

@overload
def stream_response(
    stream: Stream, 
    types: T7[P1, P2, P3, P4, P5, P6, P7]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6, P7]]:
    """Signature if exactly seven types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: T8[P1, P2, P3, P4, P5, P6, P7, P8]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6, P7, P8]]:
    """Signature if exactly eight types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: T9[P1, P2, P3, P4, P5, P6, P7, P8, P9]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6, P7, P8, P9]]:
    """Signature if exactly nine types are passed in for the `types` parameter"""
    
@overload
def stream_response(
    stream: Stream, 
    types: T10[P1, P2, P3, P4, P5, P6, P7, P8, P9, P10]
) -> Iterator[Union[P1, P2, P3, P4, P5, P6, P7, P8, P9, P10]]:
    """Signature if exactly ten types are passed in for the `types` parameter"""

# We have to be more generic in our type-hinting for the concrete implementation 
# Otherwise, MyPy struggles to figure out that it's a valid argument to `isinstance`
def stream_response(
    stream: Stream,
    types: Union[type, Tuple[type, ...]]
) -> Iterator[Packet]:
    
    while response := stream():
        if isinstance(response, Done): return
        if isinstance(response, types): yield response

def print_messages(stream: Stream) -> None:
    for m in stream_response(stream, Message):
        print(m.msg)

msgs = iter((Message(0, "hello"), Exn("Oops", (1, 42)), Done()))
print_messages(lambda: next(msgs))

【讨论】：

啊，在类型参数上使用边界的技巧非常聪明。元组重载令人不快，但至少一种情况还不错！ @Clément，同意——当注释 classmethods 返回类 ***.com/a/68283181/13990016 的实例时，单一类型的情况是一种常见模式很好的捕捉，也有很好的答案！ +1 非常感谢！ @Clément -- 我想到了一种方法，可以让这一点不那么冗长；已将其编辑到我的答案中。但是，它确实使注释的可读性更差！

以上是关于如何注释返回类型取决于其参数的函数？的主要内容，如果未能解决你的问题，请参考以下文章

当一个函数无返回值时,函数的类型应定义为啥

如何注释函数参数的最大限制

js 函数注释

如何定义具有重载的函数以恰好具有 1 或 2 个参数，这些参数的类型取决于使用的字符串文字

37 py为什么要使用函数函数中添加文本注释没有返回值的函数

如何在返回其回调之一结果的函数的 Typescript 中声明类型？