Fix type annotations for stack.py (#5566)

Author: archaengel

data_structures/stacks/balanced_parentheses.py

@@ -14,7 +14,7 @@ def balanced_parentheses(parentheses: str) -> bool:
     >>> balanced_parentheses("")
     True
     """
-    stack = Stack()
+    stack: Stack[str] = Stack()
     bracket_pairs = {"(": ")", "[": "]", "{": "}"}
     for bracket in parentheses:
         if bracket in bracket_pairs:

data_structures/stacks/dijkstras_two_stack_algorithm.py

@@ -51,8 +51,8 @@ def dijkstras_two_stack_algorithm(equation: str) -> int:
     """
     operators = {"*": op.mul, "/": op.truediv, "+": op.add, "-": op.sub}
 
-    operand_stack = Stack()
-    operator_stack = Stack()
+    operand_stack: Stack[int] = Stack()
+    operator_stack: Stack[str] = Stack()
 
     for i in equation:
         if i.isdigit():

data_structures/stacks/infix_to_postfix_conversion.py

@@ -38,7 +38,7 @@ def infix_to_postfix(expression_str: str) -> str:
     """
     if not balanced_parentheses(expression_str):
         raise ValueError("Mismatched parentheses")
-    stack = Stack()
+    stack: Stack[str] = Stack()
     postfix = []
     for char in expression_str:
         if char.isalpha() or char.isdigit():

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