python-design-patterns/behavioral/iterator_pattern.py at main · 5iing/python-design-patterns · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
from abc import *

class Iterator(ABC):
    @abstractmethod
    def __iter__(self): ...

    @abstractmethod
    def __next__(self) -> str: ...

class NameCollection:
    def __init__(self, names: list[str]) -> None:
        self.names = names

    def get_first_name_iterator(self):
        return FirstNameIterator(self.names)

    def get_last_name_iterator(self):
        return LastNameIterator(self.names)

class FirstNameIterator(Iterator):
    def __init__(self, names: list[str]) -> None:
        self.names = names
        self._index = 0

    def __iter__(self):
        return self

    def __next__(self) -> str:
        if self._index >= len(self.names):
            raise StopIteration
        name = self.names[self._index]
        self._index += 1
        return name.split()[0]

class LastNameIterator(Iterator):
    def __init__(self, names: list[str]) -> None:
        self.names = names
        self._index = 0

    def __iter__(self):
        return self

    def __next__(self) -> str:
        if self._index >= len(self.names):
            raise StopIteration
        name = self.names[self._index]
        self._index += 1
        return name.split()[-1]

if __name__ == '__main__':
    names = ["yeonwoo lee", "liam lawson", "calros sainz"]
    collection = NameCollection(names)

    for first_name in collection.get_first_name_iterator():
        print(first_name)

    print("")

    for last_name in collection.get_last_name_iterator():
        print(last_name)