Simple stack in Python

stack.py

# Author: Michelle Mackie
# Update: Ivan E. Cao-Berg
#
# Copyright (C) 2012-2025
# School of Computer Science
# Carnegie Mellon University
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published
# by the Free Software Foundation; either version 2 of the License,
# or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.
#
# For additional information visit http://www.andrew.cmu.edu/~icaoberg or
# send email to [email protected]

class Stack:
    def __init__(self):
        """
        Initializes an empty stack.
        """
        self.stack = []

    def push(self, value):
        """
        Adds a value to the top of the stack.

        Args:
            value (Any): The value to add to the stack.

        Returns:
            bool: True if the value was successfully added, False otherwise.
        """
        if value is None:
            return False

        self.stack.insert(0, value)
        return True

    def peek(self):
        """
        Returns the top value of the stack without removing it.

        Returns:
            Any: The top value of the stack, or None if the stack is empty.
        """
        return self.stack[0] if self.stack else None

    def pop(self):
        """
        Removes and returns the top value from the stack.

        Returns:
            bool: True if a value was successfully removed, False otherwise.
        """
        if self.stack:
            self.stack.pop(0)
            return True
        return False

    def size(self):
        """
        Returns the number of elements in the stack.

        Returns:
            int: The size of the stack.
        """
        return len(self.stack)

    def contains(self, value):
        """
        Checks if the stack contains a specific value.

        Args:
            value (Any): The value to check for.

        Returns:
            bool: True if the value exists in the stack, False otherwise.
        """
        return value in self.stack

    def clear(self):
        """
        Clears all elements from the stack.
        """
        self.stack.clear()

    def is_empty(self):
        """
        Checks if the stack is empty.

        Returns:
            bool: True if the stack is empty, False otherwise.
        """
        return len(self.stack) == 0

test.py

# Author: Michelle Mackie
# Update: Ivan E. Cao-Berg
#
# Copyright (C) 2012-2025
# School of Computer Science
# Carnegie Mellon University
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published
# by the Free Software Foundation; either version 2 of the License,
# or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA
# 02110-1301, USA.
#
# For additional information visit http://www.andrew.cmu.edu/~icaoberg or
# send email to [email protected]

# Simple unit tests for the Stack class
# Last tested on Python 2.7

import unittest
from stack import Stack

class TestStack(unittest.TestCase):
    def test_push(self):
        s = Stack()
        for value in range(1, 101):
            self.assertTrue(s.push(value))

    def test_pop(self):
        s = Stack()
        for value in range(1, 11):
            self.assertTrue(s.push(value))
            self.assertTrue(s.pop())

        self.assertFalse(s.pop())

    def test_peek(self):
        s = Stack()
        for value in range(1, 11):
            self.assertTrue(s.push(value))
            self.assertEqual(s.peek(), value)

    def test_contains(self):
        s = Stack()
        for value in range(1, 11):
            self.assertTrue(s.push(value))
            self.assertTrue(s.contains(value))
            self.assertFalse(s.contains(value + 1))

    def test_size(self):
        s = Stack()
        for value in range(1, 11):
            self.assertTrue(s.push(value))
            self.assertEqual(s.size(), value)

    def test_clear(self):
        s = Stack()
        for value in range(1, 11):
            self.assertTrue(s.push(value))

        self.assertEqual(s.size(), 10)
        s.clear()
        self.assertEqual(s.size(), 0)

    def test_is_empty(self):
        s = Stack()
        self.assertTrue(s.is_empty())
        s.push(1)
        self.assertFalse(s.is_empty())
        s.pop()
        self.assertTrue(s.is_empty())

if __name__ == '__main__':
    unittest.main()