This is my Stack implementation.
class Stack:
def __init__(self):
self.head = None
self.size = 0
def push(self, item):
node
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
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
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