gh-116738: Address the review comments

Author: yoney

Lib/test/test_free_threading/test_heapq.py

@@ -25,11 +25,8 @@ def test_racing_heapify(self):
         heap = list(range(OBJECT_COUNT))
         shuffle(heap)
 
-        def heapify_func(heap: list[int]):
-            heapq.heapify(heap)
-
         self.run_concurrently(
-            worker_func=heapify_func, args=(heap,), nthreads=NTHREADS
+            worker_func=heapq.heapify, args=(heap,), nthreads=NTHREADS
         )
         self.assertTrue(self.is_min_heap_property_satisfied(heap))
 
@@ -46,12 +43,10 @@ def heappush_func(heap: list[int]):
         self.assertTrue(self.is_min_heap_property_satisfied(heap))
 
     def test_racing_heappop(self):
-        heap = list(range(OBJECT_COUNT))
-        shuffle(heap)
-        heapq.heapify(heap)
+        heap = self.create_heap(OBJECT_COUNT, HeapKind.MIN)
 
         # Each thread pops (OBJECT_COUNT / NTHREADS) items
-        self.assertEqual(0, OBJECT_COUNT % NTHREADS)
+        self.assertEqual(OBJECT_COUNT % NTHREADS, 0)
         per_thread_pop_count = OBJECT_COUNT // NTHREADS
 
         def heappop_func(heap: list[int], pop_count: int):
@@ -68,16 +63,11 @@ def heappop_func(heap: list[int], pop_count: int):
             args=(heap, per_thread_pop_count),
             nthreads=NTHREADS,
         )
-        self.assertEqual(0, len(heap))
+        self.assertEqual(len(heap), 0)
 
     def test_racing_heappushpop(self):
-        heap = list(range(OBJECT_COUNT))
-        shuffle(heap)
-        heapq.heapify(heap)
-
-        pushpop_items = [
-            randint(-OBJECT_COUNT, OBJECT_COUNT) for _ in range(OBJECT_COUNT)
-        ]
+        heap = self.create_heap(OBJECT_COUNT, HeapKind.MIN)
+        pushpop_items = self.create_random_list(-5_000, 10_000, OBJECT_COUNT)
 
         def heappushpop_func(heap: list[int], pushpop_items: list[int]):
             for item in pushpop_items:
@@ -89,17 +79,12 @@ def heappushpop_func(heap: list[int], pushpop_items: list[int]):
             args=(heap, pushpop_items),
             nthreads=NTHREADS,
         )
-        self.assertEqual(OBJECT_COUNT, len(heap))
+        self.assertEqual(len(heap), OBJECT_COUNT)
         self.assertTrue(self.is_min_heap_property_satisfied(heap))
 
     def test_racing_heapreplace(self):
-        heap = list(range(OBJECT_COUNT))
-        shuffle(heap)
-        heapq.heapify(heap)
-
-        replace_items = [
-            randint(-OBJECT_COUNT, OBJECT_COUNT) for _ in range(OBJECT_COUNT)
-        ]
+        heap = self.create_heap(OBJECT_COUNT, HeapKind.MIN)
+        replace_items = self.create_random_list(-5_000, 10_000, OBJECT_COUNT)
 
         def heapreplace_func(heap: list[int], replace_items: list[int]):
             for item in replace_items:
@@ -110,18 +95,15 @@ def heapreplace_func(heap: list[int], replace_items: list[int]):
             args=(heap, replace_items),
             nthreads=NTHREADS,
         )
-        self.assertEqual(OBJECT_COUNT, len(heap))
+        self.assertEqual(len(heap), OBJECT_COUNT)
         self.assertTrue(self.is_min_heap_property_satisfied(heap))
 
     def test_racing_heapify_max(self):
         max_heap = list(range(OBJECT_COUNT))
         shuffle(max_heap)
 
-        def heapify_max_func(max_heap: list[int]):
-            heapq.heapify_max(max_heap)
-
         self.run_concurrently(
-            worker_func=heapify_max_func, args=(max_heap,), nthreads=NTHREADS
+            worker_func=heapq.heapify_max, args=(max_heap,), nthreads=NTHREADS
         )
         self.assertTrue(self.is_max_heap_property_satisfied(max_heap))
 
@@ -138,12 +120,10 @@ def heappush_max_func(max_heap: list[int]):
         self.assertTrue(self.is_max_heap_property_satisfied(max_heap))
 
     def test_racing_heappop_max(self):
-        max_heap = list(range(OBJECT_COUNT))
-        shuffle(max_heap)
-        heapq.heapify_max(max_heap)
+        max_heap = self.create_heap(OBJECT_COUNT, HeapKind.MAX)
 
         # Each thread pops (OBJECT_COUNT / NTHREADS) items
-        self.assertEqual(0, OBJECT_COUNT % NTHREADS)
+        self.assertEqual(OBJECT_COUNT % NTHREADS, 0)
         per_thread_pop_count = OBJECT_COUNT // NTHREADS
 
         def heappop_max_func(max_heap: list[int], pop_count: int):
@@ -160,16 +140,11 @@ def heappop_max_func(max_heap: list[int], pop_count: int):
             args=(max_heap, per_thread_pop_count),
             nthreads=NTHREADS,
         )
-        self.assertEqual(0, len(max_heap))
+        self.assertEqual(len(max_heap), 0)
 
     def test_racing_heappushpop_max(self):
-        max_heap = list(range(OBJECT_COUNT))
-        shuffle(max_heap)
-        heapq.heapify_max(max_heap)
-
-        pushpop_items = [
-            randint(-OBJECT_COUNT, OBJECT_COUNT) for _ in range(OBJECT_COUNT)
-        ]
+        max_heap = self.create_heap(OBJECT_COUNT, HeapKind.MAX)
+        pushpop_items = self.create_random_list(-5_000, 10_000, OBJECT_COUNT)
 
         def heappushpop_max_func(
             max_heap: list[int], pushpop_items: list[int]
@@ -183,17 +158,12 @@ def heappushpop_max_func(
             args=(max_heap, pushpop_items),
             nthreads=NTHREADS,
         )
-        self.assertEqual(OBJECT_COUNT, len(max_heap))
+        self.assertEqual(len(max_heap), OBJECT_COUNT)
         self.assertTrue(self.is_max_heap_property_satisfied(max_heap))
 
     def test_racing_heapreplace_max(self):
-        max_heap = list(range(OBJECT_COUNT))
-        shuffle(max_heap)
-        heapq.heapify_max(max_heap)
-
-        replace_items = [
-            randint(-OBJECT_COUNT, OBJECT_COUNT) for _ in range(OBJECT_COUNT)
-        ]
+        max_heap = self.create_heap(OBJECT_COUNT, HeapKind.MAX)
+        replace_items = self.create_random_list(-5_000, 10_000, OBJECT_COUNT)
 
         def heapreplace_max_func(
             max_heap: list[int], replace_items: list[int]
@@ -206,7 +176,7 @@ def heapreplace_max_func(
             args=(max_heap, replace_items),
             nthreads=NTHREADS,
         )
-        self.assertEqual(OBJECT_COUNT, len(max_heap))
+        self.assertEqual(len(max_heap), OBJECT_COUNT)
         self.assertTrue(self.is_max_heap_property_satisfied(max_heap))
 
     def is_min_heap_property_satisfied(self, heap: list[object]) -> bool:
@@ -254,25 +224,47 @@ def is_sorted_descending(lst: list[object]) -> bool:
         return all(lst[i - 1] >= lst[i] for i in range(1, len(lst)))
 
     @staticmethod
-    def run_concurrently(worker_func, args, nthreads) -> None:
+    def create_heap(size: int, heap_kind: HeapKind) -> list[int]:
+        """
+        Create a min/max heap where elements are in the range (0, size - 1) and
+        shuffled before heapify.
+        """
+        heap = list(range(OBJECT_COUNT))
+        shuffle(heap)
+        if heap_kind == HeapKind.MIN:
+            heapq.heapify(heap)
+        else:
+            heapq.heapify_max(heap)
+
+        return heap
+
+    @staticmethod
+    def create_random_list(a: int, b: int, size: int) -> list[int]:
+        """
+        Create a random list where elements are in the range a <= elem <= b
+        """
+        return [randint(-a, b) for _ in range(size)]
+
+    def run_concurrently(self, worker_func, args, nthreads) -> None:
         """
         Run the worker function concurrently in multiple threads.
         """
-        barrier = Barrier(NTHREADS)
+        barrier = Barrier(nthreads)
 
         def wrapper_func(*args):
             # Wait for all threadss to reach this point before proceeding.
             barrier.wait()
             worker_func(*args)
 
-        workers = []
-        for _ in range(nthreads):
-            worker = Thread(target=wrapper_func, args=args)
-            workers.append(worker)
-            worker.start()
+        with threading_helper.catch_threading_exception() as cm:
+            workers = (
+                Thread(target=wrapper_func, args=args) for _ in range(nthreads)
+            )
+            with threading_helper.start_threads(workers):
+                pass
 
-        for worker in workers:
-            worker.join()
+            # Worker threads should not raise any exceptions
+            self.assertIsNone(cm.exc_value)
 
 
 if __name__ == "__main__":

Misc/NEWS.d/next/Core_and_Builtins/2025-06-02-13-57-40.gh-issue-116738.ycJsL8.rst

@@ -1 +1 @@
-Make methods on :class:`heapq` thread-safe when the GIL is disabled.
+Make methods in :mod:`heapq` thread-safe on the :term:`free threaded <free threading>` build.