BST and RSA doctest (#8693)

* rsa key doctest

* move doctest to module docstring

* all tests to doctest

* moved is_right to property

* is right test

* fixed rsa doctest import

* Test error when deleting non-existing element

* fixing ruff EM102

* convert property 'is_right' to one-liner

Also use 'is' instead of '=='

Co-authored-by: Tianyi Zheng <tianyizheng02@gmail.com>

* child instead of children

Co-authored-by: Tianyi Zheng <tianyizheng02@gmail.com>

* remove type hint

* Update data_structures/binary_tree/binary_search_tree.py

---------

Co-authored-by: Tianyi Zheng <tianyizheng02@gmail.com>

ciphers/rsa_key_generator.py

@@ -2,8 +2,7 @@ import os
 import random
 import sys
 
-from . import cryptomath_module as cryptoMath  # noqa: N812
+from . import cryptomath_module, rabin_miller
-from . import rabin_miller as rabinMiller  # noqa: N812
 
 
 def main() -> None:
@@ -13,20 +12,26 @@ def main() -> None:
 
 
 def generate_key(key_size: int) -> tuple[tuple[int, int], tuple[int, int]]:
-    print("Generating prime p...")
+    """
-    p = rabinMiller.generate_large_prime(key_size)
+    >>> random.seed(0) # for repeatability
-    print("Generating prime q...")
+    >>> public_key, private_key = generate_key(8)
-    q = rabinMiller.generate_large_prime(key_size)
+    >>> public_key
+    (26569, 239)
+    >>> private_key
+    (26569, 2855)
+    """
+    p = rabin_miller.generate_large_prime(key_size)
+    q = rabin_miller.generate_large_prime(key_size)
     n = p * q
 
-    print("Generating e that is relatively prime to (p - 1) * (q - 1)...")
+    # Generate e that is relatively prime to (p - 1) * (q - 1)
     while True:
         e = random.randrange(2 ** (key_size - 1), 2 ** (key_size))
-        if cryptoMath.gcd(e, (p - 1) * (q - 1)) == 1:
+        if cryptomath_module.gcd(e, (p - 1) * (q - 1)) == 1:
             break
 
-    print("Calculating d that is mod inverse of e...")
+    # Calculate d that is mod inverse of e
-    d = cryptoMath.find_mod_inverse(e, (p - 1) * (q - 1))
+    d = cryptomath_module.find_mod_inverse(e, (p - 1) * (q - 1))
 
     public_key = (n, e)
     private_key = (n, d)

data_structures/binary_tree/binary_search_tree.py

@@ -1,5 +1,62 @@
-"""
+r"""
 A binary search Tree
+
+Example
+              8
+             / \
+            3   10
+           / \    \
+          1   6    14
+             / \   /
+            4   7 13
+
+>>> t = BinarySearchTree()
+>>> t.insert(8, 3, 6, 1, 10, 14, 13, 4, 7)
+>>> print(" ".join(repr(i.value) for i in t.traversal_tree()))
+8 3 1 6 4 7 10 14 13
+>>> print(" ".join(repr(i.value) for i in t.traversal_tree(postorder)))
+1 4 7 6 3 13 14 10 8
+>>> t.remove(20)
+Traceback (most recent call last):
+    ...
+ValueError: Value 20 not found
+>>> BinarySearchTree().search(6)
+Traceback (most recent call last):
+    ...
+IndexError: Warning: Tree is empty! please use another.
+
+Other example:
+
+>>> testlist = (8, 3, 6, 1, 10, 14, 13, 4, 7)
+>>> t = BinarySearchTree()
+>>> for i in testlist:
+...     t.insert(i)
+
+Prints all the elements of the list in order traversal
+>>> print(t)
+{'8': ({'3': (1, {'6': (4, 7)})}, {'10': (None, {'14': (13, None)})})}
+
+Test existence
+>>> t.search(6) is not None
+True
+>>> t.search(-1) is not None
+False
+
+>>> t.search(6).is_right
+True
+>>> t.search(1).is_right
+False
+
+>>> t.get_max().value
+14
+>>> t.get_min().value
+1
+>>> t.empty()
+False
+>>> for i in testlist:
+...     t.remove(i)
+>>> t.empty()
+True
 """
 
 from collections.abc import Iterable
@@ -20,6 +77,10 @@ class Node:
             return str(self.value)
         return pformat({f"{self.value}": (self.left, self.right)}, indent=1)
 
+    @property
+    def is_right(self) -> bool:
+        return self.parent is not None and self is self.parent.right
+
 
 class BinarySearchTree:
     def __init__(self, root: Node | None = None):
@@ -35,18 +96,13 @@ class BinarySearchTree:
         if new_children is not None:  # reset its kids
             new_children.parent = node.parent
         if node.parent is not None:  # reset its parent
-            if self.is_right(node):  # If it is the right children
+            if node.is_right:  # If it is the right child
                 node.parent.right = new_children
             else:
                 node.parent.left = new_children
         else:
             self.root = new_children
 
-    def is_right(self, node: Node) -> bool:
-        if node.parent and node.parent.right:
-            return node == node.parent.right
-        return False
-
     def empty(self) -> bool:
         return self.root is None
 
@@ -119,8 +175,12 @@ class BinarySearchTree:
         return node
 
     def remove(self, value: int) -> None:
-        node = self.search(value)  # Look for the node with that label
+        # Look for the node with that label
-        if node is not None:
+        node = self.search(value)
+        if node is None:
+            msg = f"Value {value} not found"
+            raise ValueError(msg)
+
         if node.left is None and node.right is None:  # If it has no children
             self.__reassign_nodes(node, None)
         elif node.left is None:  # Has only right children
@@ -128,12 +188,12 @@ class BinarySearchTree:
         elif node.right is None:  # Has only left children
             self.__reassign_nodes(node, node.left)
         else:
-                tmp_node = self.get_max(
+            predecessor = self.get_max(
                 node.left
             )  # Gets the max value of the left branch
-                self.remove(tmp_node.value)  # type: ignore
+            self.remove(predecessor.value)  # type: ignore
             node.value = (
-                    tmp_node.value  # type: ignore
+                predecessor.value  # type: ignore
             )  # Assigns the value to the node to delete and keep tree structure
 
     def preorder_traverse(self, node: Node | None) -> Iterable:
@@ -177,55 +237,6 @@ def postorder(curr_node: Node | None) -> list[Node]:
     return node_list
 
 
-def binary_search_tree() -> None:
-    r"""
-    Example
-                  8
-                 / \
-                3   10
-               / \    \
-              1   6    14
-                 / \   /
-                4   7 13
-
-    >>> t = BinarySearchTree()
-    >>> t.insert(8, 3, 6, 1, 10, 14, 13, 4, 7)
-    >>> print(" ".join(repr(i.value) for i in t.traversal_tree()))
-    8 3 1 6 4 7 10 14 13
-    >>> print(" ".join(repr(i.value) for i in t.traversal_tree(postorder)))
-    1 4 7 6 3 13 14 10 8
-    >>> BinarySearchTree().search(6)
-    Traceback (most recent call last):
-        ...
-    IndexError: Warning: Tree is empty! please use another.
-    """
-    testlist = (8, 3, 6, 1, 10, 14, 13, 4, 7)
-    t = BinarySearchTree()
-    for i in testlist:
-        t.insert(i)
-
-    # Prints all the elements of the list in order traversal
-    print(t)
-
-    if t.search(6) is not None:
-        print("The value 6 exists")
-    else:
-        print("The value 6 doesn't exist")
-
-    if t.search(-1) is not None:
-        print("The value -1 exists")
-    else:
-        print("The value -1 doesn't exist")
-
-    if not t.empty():
-        print("Max Value: ", t.get_max().value)  # type: ignore
-        print("Min Value: ", t.get_min().value)  # type: ignore
-
-    for i in testlist:
-        t.remove(i)
-        print(t)
-
-
 if __name__ == "__main__":
     import doctest