Add DocTests to diffie.py

ciphers/diffie.py

@@ -2,6 +2,25 @@
 
 
 def find_primitive(n: int) -> int | None:
+    """
+    Find a primitive root modulo n, if one exists.
+
+    Args:
+        n (int): The modulus for which to find a primitive root.
+
+    Returns:
+        int | None: The primitive root if one exists, or None if there is none.
+
+    Examples:
+        >>> find_primitive(7)  # Modulo 7 has primitive root 3
+        3
+
+        >>> find_primitive(11)  # Modulo 11 has primitive root 2
+        2
+
+        >>> find_primitive(8)  # Modulo 8 has no primitive root
+
+    """
     for r in range(1, n):
         li = []
         for x in range(n - 1):
@@ -15,18 +34,6 @@ def find_primitive(n: int) -> int | None:
 
 
 if __name__ == "__main__":
-    q = int(input("Enter a prime number q: "))
-    a = find_primitive(q)
-    if a is None:
-        print(f"Cannot find the primitive for the value: {a!r}")
-    else:
-        a_private = int(input("Enter private key of A: "))
-        a_public = pow(a, a_private, q)
-        b_private = int(input("Enter private key of B: "))
-        b_public = pow(a, b_private, q)
-
-        a_secret = pow(b_public, a_private, q)
-        b_secret = pow(a_public, b_private, q)
-
-        print("The key value generated by A is: ", a_secret)
-        print("The key value generated by B is: ", b_secret)
+    import doctest
+
+    doctest.testmod()