Fix type annotations for stack.py

Author: archaengel

data_structures/stacks/stack.py

@@ -1,5 +1,9 @@
 from __future__ import annotations
 
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
 
 class StackOverflowError(BaseException):
     pass
@@ -9,7 +13,7 @@ class StackUnderflowError(BaseException):
     pass
 
 
-class Stack:
+class Stack(Generic[T]):
     """A stack is an abstract data type that serves as a collection of
     elements with two principal operations: push() and pop(). push() adds an
     element to the top of the stack, and pop() removes an element from the top
@@ -19,7 +23,7 @@ class Stack:
     """
 
     def __init__(self, limit: int = 10):
-        self.stack: list[int] = []
+        self.stack: list[T] = []
         self.limit = limit
 
     def __bool__(self) -> bool:
@@ -28,13 +32,13 @@ def __bool__(self) -> bool:
     def __str__(self) -> str:
         return str(self.stack)
 
-    def push(self, data):
+    def push(self, data: T) -> None:
         """Push an element to the top of the stack."""
         if len(self.stack) >= self.limit:
             raise StackOverflowError
         self.stack.append(data)
 
-    def pop(self):
+    def pop(self) -> T:
         """
         Pop an element off of the top of the stack.
 
@@ -47,7 +51,7 @@ def pop(self):
             raise StackUnderflowError
         return self.stack.pop()
 
-    def peek(self):
+    def peek(self) -> T:
         """
         Peek at the top-most element of the stack.
 
@@ -71,7 +75,7 @@ def size(self) -> int:
         """Return the size of the stack."""
         return len(self.stack)
 
-    def __contains__(self, item) -> bool:
+    def __contains__(self, item: T) -> bool:
         """Check if item is in stack"""
         return item in self.stack
 
@@ -80,7 +84,7 @@ def test_stack() -> None:
     """
     >>> test_stack()
     """
-    stack = Stack(10)
+    stack: Stack[int] = Stack(10)
     assert bool(stack) is False
     assert stack.is_empty() is True
     assert stack.is_full() is False