adding static type checking to basic_binary_tree.py

data_structures/binary_tree/basic_binary_tree.py

@@ -1,59 +1,85 @@
+from typing import Optional
+
+
 class Node:
     """
-    This is the Class Node with a constructor that contains data variable to type data
-    and left, right pointers.
+    A Node has data variable and pointers to Nodes to its left and right.
     """
-
-    def __init__(self, data):
+    def __init__(self, data: int) -> None:
         self.data = data
-        self.left = None
-        self.right = None
-
+        self.left: Optional[Node] = None
+        self.right: Optional[Node] = None
 
-def display(tree):  # In Order traversal of the tree
 
-    if tree is None:
-        return
-
-    if tree.left is not None:
+def display(tree: Optional[Node]) -> None:  # In Order traversal of the tree
+    """
+    >>> root = Node(1)
+    >>> root.left = Node(0)
+    >>> root.right = Node(2)
+    >>> display(root)
+    0
+    1
+    2
+    >>> display(root.right)
+    2
+    """
+    if tree:
         display(tree.left)
-
-    print(tree.data)
-
-    if tree.right is not None:
+        print(tree.data)
         display(tree.right)
 
-    return
-
 
-def depth_of_tree(
-    tree,
-):  # This is the recursive function to find the depth of binary tree.
-    if tree is None:
-        return 0
-    else:
-        depth_l_tree = depth_of_tree(tree.left)
-        depth_r_tree = depth_of_tree(tree.right)
-        if depth_l_tree > depth_r_tree:
-            return 1 + depth_l_tree
-        else:
-            return 1 + depth_r_tree
+def depth_of_tree(tree: Optional[Node]) -> int:
+    """
+    Recursive function that returns the depth of a binary tree.
+
+    >>> root = Node(0)
+    >>> depth_of_tree(root)
+    1
+    >>> root.left = Node(0)
+    >>> depth_of_tree(root)
+    2
+    >>> root.right = Node(0)
+    >>> depth_of_tree(root)
+    2
+    >>> root.left.right = Node(0)
+    >>> depth_of_tree(root)
+    3
+    >>> depth_of_tree(root.left)
+    2
+    """
+    return 1 + max(depth_of_tree(tree.left), depth_of_tree(tree.right)) if tree else 0
 
 
-def is_full_binary_tree(
-    tree,
-):  # This function returns that is it full binary tree or not?
-    if tree is None:
-        return True
-    if (tree.left is None) and (tree.right is None):
+def is_full_binary_tree(tree: Node) -> bool:
+    """
+    Returns True if this is a full binary tree
+
+    >>> root = Node(0)
+    >>> is_full_binary_tree(root)
+    True
+    >>> root.left = Node(0)
+    >>> is_full_binary_tree(root)
+    False
+    >>> root.right = Node(0)
+    >>> is_full_binary_tree(root)
+    True
+    >>> root.left.left = Node(0)
+    >>> is_full_binary_tree(root)
+    False
+    >>> root.right.right = Node(0)
+    >>> is_full_binary_tree(root)
+    False
+    """
+    if not tree:
         return True
-    if (tree.left is not None) and (tree.right is not None):
+    if tree.left and tree.right:
         return is_full_binary_tree(tree.left) and is_full_binary_tree(tree.right)
     else:
-        return False
+        return not tree.left and not tree.right
 
 
-def main():  # Main function for testing.
+def main() -> None:  # Main function for testing.
     tree = Node(1)
     tree.left = Node(2)
     tree.right = Node(3)