Segmented sieve - doctests

maths/segmented_sieve.py

@@ -4,7 +4,26 @@
 
 
 def sieve(n: int) -> list[int]:
-    """Segmented Sieve."""
+    """
+    Segmented Sieve.
+
+    Examples:
+    >>> sieve(8)
+    [2, 3, 5, 7]
+
+    >>> sieve(27)
+    [2, 3, 5, 7, 11, 13, 17, 19, 23]
+
+    """
+
+    if n < 0:
+        word = f"Number should not be negative n {n}"
+        raise ValueError(word)
+
+    if n == 0:
+        word = f"Number should not be zero n {n}"
+        raise ValueError(word)
+
     in_prime = []
     start = 2
     end = int(math.sqrt(n))  # Size of every segment
@@ -42,4 +61,9 @@ def sieve(n: int) -> list[int]:
     return prime
 
 
-print(sieve(10**6))
+if __name__ == "__main__":
+    import doctest
+
+    doctest.testmod()
+
+    print(sieve(10**6))