python给__iter__添加type hint类型注解

时间:2024-06-01 14:06:33

Summary::

from typing import Iterator

class XXX:
    def __iter__(self) -> Iterator:
        ...

Example::

from typing import Iterable, Iterator, List, Sized


class Scoreboard1(Iterable):

   def __iter__(self) -> Iterator:
      """This method should return an Iterator.
         Code that wants an iterator will call iter(obj)
         to create one.
      """
      return iter(())


class Scoreboard2(Iterable[int]):
   def __init__(self) -> None:
       self.data: List[int] = []

   def __iter__(self) -> Iterator[int]:
      """This method should return an Iterator.
         Code that wants an iterator will call iter(obj)
         to create one.
      """
      # create an iterator from a range
      return iter(self.data)


class Course:
    pass


class CourseList(Sized):
    """List of courses taken by a student"""

    def __init__(self, student_id) -> None:
        self.student_id = student_id
        self.courses: 'List[Course]' = []

    def add_course(self, course: 'Course') -> None:
        self.courses.append(course)

    def __len__(self) -> int:
        return len(self.courses)


class CourseList2(CourseList):
   def __iter__(self) -> 'Iterator[Course]':
       return iter(self.courses)


def _test() -> None:
    assert list(Scoreboard1()) == []
    sb2 = Scoreboard2()
    assert list(sb2) == []
    sb2.data.append(1)
    assert list(sb2) == [1]
    try:
        list(sb2)
    except TypeError as e:
        assert str(e) == "TypeError: object of type 'Scoreboard2' has no len()"
    cs = CourseList(1)
    assert len(cs) == 0
    try:
        list(cs)  # type:ignore
    except TypeError as err:
        assert str(err) == "'CourseList' object is not iterable"
    cs2 = CourseList2(1)
    assert len(cs2) == 0
    assert list(cs2) == []
    print('Pass all tests~')


if __name__ == '__main__':
    _test()

Ref::

Type Hints –– An Introduction | Individual Software Process