From 9a03919052276bfe83f90a7ba2fe258d20a104e9 Mon Sep 17 00:00:00 2001
From: Edward Nuno <GOD11341258@gmail.com>
Date: Tue, 26 Oct 2021 12:12:46 -0700
Subject: [PATCH] [mypy] Fix type annotations for stack_using_dll.py (#5577)

* Fix mypy annotations for stack_using_dll.py

* Replace Optional with inline union type
---
 data_structures/stacks/stack_using_dll.py | 40 ++++++++++++++---------
 1 file changed, 24 insertions(+), 16 deletions(-)

diff --git a/data_structures/stacks/stack_using_dll.py b/data_structures/stacks/stack_using_dll.py
index 75e0cd206..a129665f2 100644
--- a/data_structures/stacks/stack_using_dll.py
+++ b/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,21 +55,23 @@ class Stack:
             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:
@@ -71,10 +79,10 @@ class Stack:
             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 @@ class Stack:
 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")