6.2. Protocol Iterator

6.2.1. Rationale

  • Used for iterating in a for loop

6.2.2. Protocol

  • __iter__(self) -> self

  • __next__(self) -> raise StopIteration

  • iter(obj) -> obj.__iter__()

  • next(obj) -> obj.__next__()

class Iterator:
    def __iter__(self):
        self._current = 0
        return self

    def __next__(self):
        if self._current >= len(self.values):
            raise StopIteration
        element = self.values[self._current]
        self._current += 1
        return element

6.2.3. Example

class Crew:
    def __init__(self):
        self.members = list()

    def __iadd__(self, other):
        self.members.append(other)
        return self

    def __iter__(self):
        self._current = 0
        return self

    def __next__(self):
        if self._current >= len(self.members):
            raise StopIteration
        result = self.members[self._current]
        self._current += 1
        return result


crew = Crew()
crew += 'Mark Watney'
crew += 'Jose Jimenez'
crew += 'Melissa Lewis'

for member in crew:
    print(member)

# Mark Watney
# Jose Jimenez
# Melissa Lewis

6.2.4. Loop and Iterators

For loop:

DATA = [1, 2, 3]

for current in DATA:
    print(current)

Intuitive implementation of the for loop:

DATA = [1, 2, 3]
iterator = iter(DATA)

try:
    current = next(iterator)
    print(current)

    current = next(iterator)
    print(current)

    current = next(iterator)
    print(current)

    current = next(iterator)
    print(current)
except StopIteration:
    pass

Intuitive implementation of the for loop:

DATA = [1, 2, 3]
iterator = DATA.__iter__()

try:
    current = iterator.__next__()
    print(current)

    current = iterator.__next__()
    print(current)

    current = iterator.__next__()
    print(current)

    current = iterator.__next__()
    print(current)
except StopIteration:
    pass

6.2.5. Built-in Type Iteration

Iterating str:

for character in 'hello':
    print(character)

# h
# e
# l
# l
# o

Iterating sequences:

for number in [1, 2, 3]:
    print(number)

# 1
# 2
# 3

Iterating dict:

DATA = {'a': 1, 'b': 2, 'c': 3}

for element in DATA:
    print(element)

# a
# b
# c

Iterating dict:

for key, value in DATA.items():
    print(f'{key} -> {value}')

# a -> 1
# b -> 2
# c -> 3

Iterating nested sequences:

for key, value in [('a',1), ('b',2), ('c',3)]:
    print(f'{key} -> {value}')

# a -> 1
# b -> 2
# c -> 3

6.2.6. Use Cases

Iterator implementation:

class Parking:
    def __init__(self):
        self._parked_cars = list()

    def park(self, car):
        self._parked_cars.append(car)

    def __iter__(self):
        self._current = 0
        return self

    def __next__(self):
        if self._current >= len(self._parked_cars):
            raise StopIteration
        element = self._parked_cars[self._current]
        self._current += 1
        return element


parking = Parking()
parking.park('Mercedes')
parking.park('Maluch')
parking.park('Toyota')

for car in parking:
    print(car)

# Mercedes
# Maluch
# Toyota

6.2.7. Assignments

Code 6.1. Solution
"""
* Assignment: Protocol Iterator Implementation
* Complexity: easy
* Lines of code: 14 lines
* Time: 5 min

English:
    1. Modify classes to implement iterator protocol
    2. Iterator should return instances of `Mission`
    3. All tests must pass
    4. Run doctests - all must succeed

Polish:
    1. Zmodyfikuj klasy aby zaimplementować protokół iterator
    2. Iterator powinien zwracać instancje `Mission`
    3. Wszystkie testy muszą przejść
    4. Uruchom doctesty - wszystkie muszą się powieść

Tests:
    >>> import sys; sys.tracebacklimit = 0
    >>> from inspect import isclass, ismethod

    >>> assert isclass(Astronaut)

    >>> astro = Astronaut('Mark', 'Watney')
    >>> assert hasattr(astro, 'firstname')
    >>> assert hasattr(astro, 'lastname')
    >>> assert hasattr(astro, 'missions')
    >>> assert hasattr(astro, '__iter__')
    >>> assert hasattr(astro, '__next__')
    >>> assert ismethod(astro.__iter__)
    >>> assert ismethod(astro.__next__)

    >>> astro = Astronaut('Jan', 'Twardowski', missions=(
    ...     Mission(1969, 'Apollo 11'),
    ...     Mission(2024, 'Artemis 3'),
    ...     Mission(2035, 'Ares 3'),
    ... ))

    >>> for mission in astro:
    ...     print(mission)
    Mission(year=1969, name='Apollo 11')
    Mission(year=2024, name='Artemis 3')
    Mission(year=2035, name='Ares 3')
"""

from dataclasses import dataclass


@dataclass
class Astronaut:
    firstname: str
    lastname: str
    missions: tuple = ()


@dataclass
class Mission:
    year: int
    name: str


Code 6.2. Solution
"""
* Assignment: Protocol Iterator Range
* Complexity: medium
* Lines of code: 14 lines
* Time: 8 min

English:
    1. Define class `Range` with parameters: `start`, `stop`, `step`
    2. Write own implementation of a built-in `range(start, stop, step)` function
    3. Assume, that user will never giv only one argument; always it will be either two or three arguments
    4. Use Iterator protocol
    5. All tests must pass
    6. Run doctests - all must succeed

Polish:
    1. Zdefiniuj klasę `Range` z parametrami: `start`, `stop`, `step`
    2. Zaimplementuj własne rozwiązanie wbudowanej funkcji `range(start, stop, step)`
    3. Przyjmij, że użytkownik nigdy nie poda tylko jednego argumentu; zawsze będą to dwa lub trzy argumenty
    4. Użyj protokołu Iterator
    5. Wszystkie testy muszą przejść
    6. Uruchom doctesty - wszystkie muszą się powieść

Tests:
    >>> import sys; sys.tracebacklimit = 0
    >>> from inspect import isclass, ismethod

    >>> assert isclass(Range)

    >>> r = Range(0, 0, 0)
    >>> assert hasattr(r, '__iter__')
    >>> assert hasattr(r, '__next__')
    >>> assert ismethod(r.__iter__)
    >>> assert ismethod(r.__next__)

    >>> list(Range(0, 10, 2))
    [0, 2, 4, 6, 8]

    >>> list(Range(0, 5))
    [0, 1, 2, 3, 4]
"""