stdlib/heapq.codon

stdlib/heapq.codon

@@ -1,9 +1,11 @@
+# (c) 2022 Exaloop Inc. All rights reserved.
+
 # TODO: heapq.merge
 
 # 'heap' is a heap at all indices >= startpos, except possibly for pos.  pos
 # is the index of a leaf with a possibly out-of-order value.  Restore the
 # heap invariant.
-def _siftdown[T](heap: List[T], startpos: int, pos: int):
+def _siftdown(heap: List[T], startpos: int, pos: int, T: type) -> void:
     newitem = heap[pos]
     # Follow the path to the root, moving parents down until finding a place
     # newitem fits.
@@ -17,7 +19,8 @@ def _siftdown[T](heap: List[T], startpos: int, pos: int):
         break
     heap[pos] = newitem
 
-def _siftup[T](heap: List[T], pos: int):
+
+def _siftup(heap: List[T], pos: int, T: type) -> void:
     endpos = len(heap)
     startpos = pos
     newitem = heap[pos]
@@ -37,8 +40,9 @@ def _siftup[T](heap: List[T], pos: int):
     heap[pos] = newitem
     _siftdown(heap, startpos, pos)
 
-def _siftdown_max[T](heap: List[T], startpos: int, pos: int):
-    'Maxheap variant of _siftdown'
+
+def _siftdown_max(heap: List[T], startpos: int, pos: int, T: type) -> void:
+    "Maxheap variant of _siftdown"
     newitem = heap[pos]
     # Follow the path to the root, moving parents down until finding a place
     # newitem fits.
@@ -52,8 +56,9 @@ def _siftdown_max[T](heap: List[T], startpos: int, pos: int):
         break
     heap[pos] = newitem
 
-def _siftup_max[T](heap: List[T], pos: int):
-    'Maxheap variant of _siftup'
+
+def _siftup_max(heap: List[T], pos: int, T: type) -> void:
+    "Maxheap variant of _siftup"
     endpos = len(heap)
     startpos = pos
     newitem = heap[pos]
@@ -73,12 +78,14 @@ def _siftup_max[T](heap: List[T], pos: int):
     heap[pos] = newitem
     _siftdown_max(heap, startpos, pos)
 
-def heappush[T](heap: List[T], item: T):
+
+def heappush(heap: List[T], item: T, T: type) -> void:
     """Push item onto heap, maintaining the heap invariant."""
     heap.append(item)
     _siftdown(heap, 0, len(heap) - 1)
 
-def heappop[T](heap: List[T]):
+
+def heappop(heap: List[T], T: type) -> T:
     """Pop the smallest item off the heap, maintaining the heap invariant."""
     lastelt = heap.pop()  # raises appropriate IndexError if heap is empty
     if heap:
@@ -88,7 +95,8 @@ def heappop[T](heap: List[T]):
         return returnitem
     return lastelt
 
-def heapreplace[T](heap: List[T], item: T):
+
+def heapreplace(heap: List[T], item: T, T: type) -> T:
     """
     Pop and return the current smallest value, and add the new item.
     This is more efficient than heappop() followed by heappush(), and can be
@@ -102,14 +110,16 @@ def heapreplace[T](heap: List[T], item: T):
     _siftup(heap, 0)
     return returnitem
 
-def heappushpop[T](heap: List[T], item: T):
+
+def heappushpop(heap: List[T], item: T, T: type) -> T:
     """Fast version of a heappush followed by a heappop."""
     if heap and heap[0] < item:
         item, heap[0] = heap[0], item
         _siftup(heap, 0)
     return item
 
-def heapify[T](x: List[T]):
+
+def heapify(x: List[T], T: type) -> void:
     """Transform list into a heap, in-place, in $O(len(x))$ time."""
     n = len(x)
     # Transform bottom-up.  The largest index there's any point to looking at
@@ -120,7 +130,8 @@ def heapify[T](x: List[T]):
     for i in reversed(range(n // 2)):
         _siftup(x, i)
 
-def _heappop_max[T](heap: List[T]):
+
+def _heappop_max(heap: List[T], T: type) -> T:
     """Maxheap version of a heappop."""
     lastelt = heap.pop()  # raises appropriate IndexError if heap is empty
     if heap:
@@ -130,20 +141,23 @@ def _heappop_max[T](heap: List[T]):
         return returnitem
     return lastelt
 
-def _heapreplace_max[T](heap: List[T], item: T):
+
+def _heapreplace_max(heap: List[T], item: T, T: type) -> T:
     """Maxheap version of a heappop followed by a heappush."""
     returnitem = heap[0]  # raises appropriate IndexError if heap is empty
     heap[0] = item
     _siftup_max(heap, 0)
     return returnitem
 
-def _heapify_max[T](x: List[T]):
+
+def _heapify_max(x: List[T], T: type) -> void:
     """Transform list into a maxheap, in-place, in O(len(x)) time."""
     n = len(x)
     for i in reversed(range(n // 2)):
         _siftup_max(x, i)
 
-def nsmallest[T](n: int, iterable: Generator[T], key = Optional[int]()):
+
+def nsmallest(n: int, iterable: Generator[T], key=Optional[int](), T: type) -> List[T]:
     """Find the n smallest elements in a dataset.
     Equivalent to:  sorted(iterable, key=key)[:n]
     """
@@ -217,7 +231,8 @@ def nsmallest[T](n: int, iterable: Generator[T], key = Optional[int]()):
         result.sort()
         return [elem for k, order, elem in result]
 
-def nlargest[T](n: int, iterable: Generator[T], key = Optional[int]()):
+
+def nlargest(n: int, iterable: Generator[T], key=Optional[int](), T: type) -> List[T]:
     """Find the n largest elements in a dataset.
     Equivalent to:  sorted(iterable, key=key, reverse=True)[:n]
     """