Destructive Stack Iteration

Question

This is my Stack implementation.

class Stack:
    def __init__(self):
        self.head = None
        self.size = 0

    def push(self, item):
        node = Node(item)
        if not self.head:
            self.head = node
        else:
            node.next = self.head
            self.head = node
        self.size += 1

    def pop(self):
        if self.size == 0:
            raise ValueError('Popping off an empty stack!')
        item = self.head.val
        self.head = self.head.next
        return item

    def peek(self):
        if self.size == 0:
            raise ValueError('Peeking into an empty stack!')
        return self.head.val

    def __iter__(self):
        return self

    def __next__(self):
        if self.head:
            curr = self.head
        else:
            raise StopIteration()
        self.head = self.head.next
        return curr.val

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


if __name__ == '__main__':
    stack = Stack()
    stack.push(12)
    stack.push(13)
    stack.push(9)
    for item in stack:
        print(item)
    print(stack.peek())

The problem with this is the iteration. The iteration is destructive and thus the call to peek at the end of the iteration throws an error.

return self.head.val AttributeError: 'NoneType' object has no attribute 'val' How do I make the iteration non-desctructive?

Answer 1

Because you can have more than one iterator on the same stack, __iter__ has to return a iterator-object, that iterates the stack:

class Stack:
    def __init__(self):
        self.head = None
        self.size = 0

    def push(self, item):
        node = Node(item)
        if self.head:
            node.next = self.head
        self.head = node
        self.size += 1

    def pop(self):
        if self.size == 0:
            raise ValueError('Popping off an empty stack!')
        item = self.head.val
        self.head = self.head.next
        return item

    def peek(self):
        if self.size == 0:
            raise ValueError('Peeking into an empty stack!')
        return self.head.val

    def __iter__(self):
        return StackIterator(self.head)

class StackIterator:
    def __init__(self, head):
        self.head = head

    def __iter__(self):
        return self

    def __next__(self):
        if not self.head:
            raise StopIteration()
        item = self.head.val
        self.head = self.head.next
        return item

Answer 2

The problem is that you do not make a difference between the iterable Stack itself and the iterator that __iter__ should return. __next__ should be called on said iterator, not on the Stack itself.

I propose the following solution:

class StackIterator:
    def __init__(self, stack):
        self.head = stack.head

    def __iter__(self):
        return self

    def __next__(self):
        if not self.head:
            raise StopIteration

        current = self.head
        self.head = self.head.next

        return current.val

Get rid of __next__ in Stack and adjust __iter__ to:

def __iter__(self):
    return StackIterator(self)

Demo:

>>> stack = Stack()
>>> stack.push(12)
>>> stack.push(13)
>>> stack.push(9)
>>> 
>>> for x in stack:
...     print(x)
... 
9
13
12
>>> stack.peek()
9

Answer 3

A simple way to do this to give your Stack an alternative head that it can use for iterating. I've also added a __len__ method tha returns the Stack's size.

class Stack:
    def __init__(self):
        self.head = None
        self.size = 0

    def __len__(self):
        return self.size

    def push(self, item):
        node = Node(item)
        if not self.head:
            self.head = node
        else:
            node.next = self.head
            self.head = node
        self.size += 1

    def pop(self):
        if self.size == 0:
            raise ValueError('Popping off an empty stack!')
        item = self.head.val
        self.head = self.head.next
        return item

    def peek(self):
        if self.size == 0:
            raise ValueError('Peeking into an empty stack!')
        return self.head.val

    def __iter__(self):
        self.top = self.head
        return self

    def __next__(self):
        if self.top:
            curr = self.top
        else:
            raise StopIteration()
        self.top = self.top.next
        return curr.val

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

if __name__ == '__main__':
    stack = Stack()
    stack.push(12)
    stack.push(13)
    stack.push(9)
    print('Size', len(stack))
    for item in stack:
        print(item)
    print(stack.peek())
    stack.push(42)
    print('Size', len(stack))
    for item in stack:
        print(item)

output

Size 3
9
13
12
9
Size 4
42
9
13
12

It's probably a Good Idea to add self.top = None to __init__, although it's not strictly necessary. And you may wish to mark it as a private name with a leading underscore: self._top.

As timgeb implies in the comments, this approach is a little fragile, since we can only perform one iteration at a time on the stack, since we only have a single self.top.

---

BTW, you can optimize that push method slightly:

def push(self, item):
    node = Node(item)
    if self.head:
        node.next = self.head
    self.head = node
    self.size += 1