Misc Itertools recipe tweaks (GH-100493)

1 files changed, 49 insertions, 5 deletions

diff --git a/Doc/library/itertools.rst b/Doc/library/itertools.rst
index 9688aae..b3634ae 100644
--- a/Doc/library/itertools.rst
+++ b/Doc/library/itertools.rst
@@ -788,6 +788,11 @@ which incur interpreter overhead.
 
 .. testcode::
 
+   import collections
+   import math
+   import operator
+   import random
+
    def take(n, iterable):
        "Return first n items of the iterable as a list"
        return list(islice(iterable, n))
@@ -892,6 +897,21 @@ which incur interpreter overhead.
        data[2] = 1
        return iter_index(data, 1) if n > 2 else iter([])
 
+   def factor(n):
+       "Prime factors of n."
+       # factor(97) --> 97
+       # factor(98) --> 2 7 7
+       # factor(99) --> 3 3 11
+       for prime in sieve(n+1):
+           while True:
+               quotient, remainder = divmod(n, prime)
+               if remainder:
+                   break
+               yield prime
+               n = quotient
+               if n == 1:
+                   return
+
    def flatten(list_of_lists):
        "Flatten one level of nesting"
        return chain.from_iterable(list_of_lists)
@@ -1134,11 +1154,6 @@ which incur interpreter overhead.
 
     Now, we test all of the itertool recipes
 
-    >>> import operator
-    >>> import collections
-    >>> import math
-    >>> import random
-
     >>> take(10, count())
     [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
 
@@ -1251,6 +1266,35 @@ which incur interpreter overhead.
     >>> set(sieve(10_000)).isdisjoint(carmichael)
     True
 
+    list(factor(0))
+    []
+    list(factor(1))
+    []
+    list(factor(2))
+    [2]
+    list(factor(3))
+    [3]
+    list(factor(4))
+    [2, 2]
+    list(factor(5))
+    [5]
+    list(factor(6))
+    [2, 3]
+    list(factor(7))
+    [7]
+    list(factor(8))
+    [2, 2, 2]
+    list(factor(9))
+    [3, 3]
+    list(factor(10))
+    [2, 5]
+    all(math.prod(factor(n)) == n for n in range(1, 1000))
+    True
+    all(set(factor(n)) <= set(sieve(n+1)) for n in range(1, 1000))
+    True
+    all(list(factor(n)) == sorted(factor(n)) for n in range(1, 1000))
+    True
+
     >>> list(flatten([('a', 'b'), (), ('c', 'd', 'e'), ('f',), ('g', 'h', 'i')]))
     ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i']