Итератор класса Python

Question

У меня есть класс узла, который содержит его родителя, и я хочу создать на нем итератор. Вот моя попытка:

class Node:
    def __init__(self, parent=None):
        self._parent = parent

    def __iter__(self):
        self = self.parent

    def __next__(self):
        if self.parent is None:
            raise StopIteration
        else:
            self = self.parent
            return self

Но когда я пытаюсь перебрать экземпляр, он никогда не останавливается и возвращает одно и то же значение, что я сделал не так?

Answer 1

Причина, по которой ваш код не работает, заключается в том, что вы пытаетесь отслеживать текущий узел в итераторе, присваивая self, который является просто локальной переменной, поэтому ничего не обновляется.

Правильный способ - извлечь класс итератора и отследить текущий узел:

class Node:
    def __init__(self, parent=None):
        self.parent = parent

    def __iter__(self):
        return NodeIterator(self)


class NodeIterator:
    def __init__(self, node):
        self.next_node = node

    def __iter__(self):
        return self

    def __next__(self):
        if self.next_node is None:
            raise StopIteration
        else:
            current_node = self.next_node
            self.next_node = self.next_node.parent
            return current_node

Это можно использовать так:

root = Node()

inner_1 = Node(root)
leaf_1 = Node(inner_1)

inner_2 = Node(root)
inner_2_1 = Node(inner_2)
leaf_2 = Node(inner_2_1)

for node in leaf_2:
    # will loop through:
    # leaf_2,
    # inner_2_1;
    # inner_2,
    # root