Generic Iterator annotations Python-CodePudding

I am trying to annotate an Iterator which only returns two values, T and cls[T].

Currently I have it annotated like this:

from __future__ import annotations
import typing


class Node(typing.Generic[T]):
    def __init__(self, value: T, next: typing.Optional[Node[T]] = None) -> None:
        self.value = value 
        self.next = next

    def __iter__(self) -> typing.Iterator[typing.Union[T, Node[T]]]:
        yield from (self.value, self.next)

This code works without any errors. However, I wish to take advantage of the fact that this dunder method will always yield two values to simplify the process for the user.

As it is, the user has to deal with the following:

one = Node[int](1, Node[int](2))
value, next = one # value = 1, next = Node(2) (i.e one.next)

# This is a typing error because next can be either int or Node[int] 
# and 'two' is expected to be Node[int]:
two: Node[int] = next

# we can fix this by doing type narrowing
assert isinstance(next, Node)
two = next # now the error should be gone

So basically I want to take advantage of the fact that the second thing returned by the __iter__ is always of type Node[T] to avoid having to do type narrowing.

I know that I have to change typing.Union[T, Node[T]] in the method's return annotation, but I do not know what to change it to.

CodePudding user response：

This is not possible to annotate using the generic Iterator. It is a class expecting exactly one type argument. (current typeshed source) That means every value returned by its __next__ method necessarily has the same type.

You are invoking the iterator protocol on that tuple of self.value, self.next. A tuple has an arbitrary number of type arguments (see here), but an iterator over it must still have exactly one. This actually leads to typing issues fairly often.

Since you seem to intend your Node class to essentially emulate the tuple interface, this may be one of the rare cases, where it is better to inherit from it directly. A tuple will obviously also give you the iterable protocol, so you can still unpack it as before, but the types should be inferred properly, if you do everything correctly.

Here is a full working example:

from __future__ import annotations
from typing import TypeVar, Optional


T = TypeVar("T")


class Node(tuple[T, "Node[T]"]):
    def __new__(cls, value: T, next_: Optional[Node[T]] = None) -> Node[T]:
        return tuple.__new__(cls, (value, next_))

    def __init__(self, value: T, next_: Optional[Node[T]] = None) -> None:
        self.value = value
        self.next = next_


if __name__ == "__main__":
    node_1 = Node(1, Node(2))
    val: int
    node_2: Node[int]
    val, node_2 = node_1

This passes mypy --strict without problems.

As an unrelated side note, I would advise against using built-in names like next.

Also, note that you do not need to specify the type argument for Node, when you initialize one because it is automatically bound to the type passed to the value parameter.

CodePudding user response：

for while are the loops that can be used and mentioned as generic iterator annotations