[mypy] Fix type annotations for stack_using_dll.py

data_structures/stacks/stack_using_dll.py

@@ -1,15 +1,21 @@
 # A complete working Python program to demonstrate all
 # stack operations using a doubly linked list
 
+from __future__ import annotations
 
-class Node:
-    def __init__(self, data):
+from typing import Generic, TypeVar
+
+T = TypeVar("T")
+
+
+class Node(Generic[T]):
+    def __init__(self, data: T):
         self.data = data  # Assign data
-        self.next = None  # Initialize next as null
-        self.prev = None  # Initialize prev as null
+        self.next: Node[T] | None = None  # Initialize next as null
+        self.prev: Node[T] | None = None  # Initialize prev as null
 
 
-class Stack:
+class Stack(Generic[T]):
     """
     >>> stack = Stack()
     >>> stack.is_empty()
@@ -35,10 +41,10 @@ class Stack:
     2->1->0->
     """
 
-    def __init__(self):
-        self.head = None
+    def __init__(self) -> None:
+        self.head: Node[T] | None = None
 
-    def push(self, data):
+    def push(self, data: T) -> None:
         """add a Node to the stack"""
         if self.head is None:
             self.head = Node(data)
@@ -49,32 +55,34 @@ def push(self, data):
             new_node.prev = None
             self.head = new_node
 
-    def pop(self):
+    def pop(self) -> T | None:
         """pop the top element off the stack"""
         if self.head is None:
             return None
         else:
+            assert self.head is not None
             temp = self.head.data
             self.head = self.head.next
-            self.head.prev = None
+            if self.head is not None:
+                self.head.prev = None
             return temp
 
-    def top(self):
+    def top(self) -> T | None:
         """return the top element of the stack"""
-        return self.head.data
+        return self.head.data if self.head is not None else None
 
-    def __len__(self):
+    def __len__(self) -> int:
         temp = self.head
         count = 0
         while temp is not None:
             count += 1
             temp = temp.next
         return count
 
-    def is_empty(self):
+    def is_empty(self) -> bool:
         return self.head is None
 
-    def print_stack(self):
+    def print_stack(self) -> None:
         print("stack elements are:")
         temp = self.head
         while temp is not None:
@@ -86,7 +94,7 @@ def print_stack(self):
 if __name__ == "__main__":
 
     # Start with the empty stack
-    stack = Stack()
+    stack: Stack[int] = Stack()
 
     # Insert 4 at the beginning. So stack becomes 4->None
     print("Stack operations using Doubly LinkedList")