ai-content-maker/.venv/Lib/site-packages/numba/tests/test_heapq.py

485 lines
14 KiB
Python

import heapq as hq
import itertools
import numpy as np
from numba import jit, typed
from numba.tests.support import TestCase, MemoryLeakMixin
def heapify(x):
return hq.heapify(x)
def heappop(heap):
return hq.heappop(heap)
def heappush(heap, item):
return hq.heappush(heap, item)
def heappushpop(heap, item):
return hq.heappushpop(heap, item)
def heapreplace(heap, item):
return hq.heapreplace(heap, item)
def nsmallest(n, iterable):
return hq.nsmallest(n, iterable)
def nlargest(n, iterable):
return hq.nlargest(n, iterable)
class _TestHeapq(MemoryLeakMixin):
def setUp(self):
super(_TestHeapq, self).setUp()
self.rnd = np.random.RandomState(42)
def test_heapify_basic_sanity(self):
pyfunc = heapify
cfunc = jit(nopython=True)(pyfunc)
a = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
b = self.listimpl(a)
pyfunc(a)
cfunc(b)
self.assertPreciseEqual(a, list(b))
# includes non-finite elements
element_pool = [3.142, -10.0, 5.5, np.nan, -np.inf, np.inf]
# list which may contain duplicate elements
for x in itertools.combinations_with_replacement(element_pool, 6):
a = list(x)
b = self.listimpl(a)
pyfunc(a)
cfunc(b)
self.assertPreciseEqual(a, list(b))
# single element list
for i in range(len(element_pool)):
a = [element_pool[i]]
b = self.listimpl(a)
pyfunc(a)
cfunc(b)
self.assertPreciseEqual(a, list(b))
# elements are tuples
a = [(3, 33), (1, 11), (2, 22)]
b = self.listimpl(a)
pyfunc(a)
cfunc(b)
self.assertPreciseEqual(a, list(b))
def check_invariant(self, heap):
for pos, item in enumerate(heap):
if pos:
parentpos = (pos - 1) >> 1
self.assertTrue(heap[parentpos] <= item)
def test_push_pop(self):
# inspired by
# https://github.com/python/cpython/blob/e42b7051/Lib/test/test_heapq.py
pyfunc_heappush = heappush
cfunc_heappush = jit(nopython=True)(pyfunc_heappush)
pyfunc_heappop = heappop
cfunc_heappop = jit(nopython=True)(pyfunc_heappop)
heap = self.listimpl([-1.0])
data = self.listimpl([-1.0])
self.check_invariant(heap)
for i in range(256):
item = self.rnd.randn(1).item(0)
data.append(item)
cfunc_heappush(heap, item)
self.check_invariant(heap)
results = []
while heap:
item = cfunc_heappop(heap)
self.check_invariant(heap)
results.append(item)
data_sorted = data[:]
data_sorted.sort()
self.assertPreciseEqual(list(data_sorted), results)
self.check_invariant(results)
def test_heapify(self):
# inspired by
# https://github.com/python/cpython/blob/e42b7051/Lib/test/test_heapq.py
pyfunc = heapify
cfunc = jit(nopython=True)(pyfunc)
for size in list(range(1, 30)) + [20000]:
heap = self.listimpl(self.rnd.random_sample(size))
cfunc(heap)
self.check_invariant(heap)
def test_heapify_exceptions(self):
pyfunc = heapify
cfunc = jit(nopython=True)(pyfunc)
# Exceptions leak references
self.disable_leak_check()
with self.assertTypingError() as e:
cfunc((1, 5, 4))
msg = 'heap argument must be a list'
self.assertIn(msg, str(e.exception))
with self.assertTypingError() as e:
cfunc(self.listimpl([1 + 1j, 2 - 3j]))
msg = ("'<' not supported between instances "
"of 'complex' and 'complex'")
self.assertIn(msg, str(e.exception))
def test_heappop_basic_sanity(self):
pyfunc = heappop
cfunc = jit(nopython=True)(pyfunc)
def a_variations():
yield [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
yield [(3, 33), (1, 111), (2, 2222)]
yield np.full(5, fill_value=np.nan).tolist()
yield np.linspace(-10, -5, 100).tolist()
for a in a_variations():
heapify(a)
b = self.listimpl(a)
for i in range(len(a)):
val_py = pyfunc(a)
val_c = cfunc(b)
self.assertPreciseEqual(a, list(b))
self.assertPreciseEqual(val_py, val_c)
def test_heappop_exceptions(self):
pyfunc = heappop
cfunc = jit(nopython=True)(pyfunc)
# Exceptions leak references
self.disable_leak_check()
with self.assertTypingError() as e:
cfunc((1, 5, 4))
msg = 'heap argument must be a list'
self.assertIn(msg, str(e.exception))
def iterables(self):
yield self.listimpl([1, 3, 5, 7, 9, 2, 4, 6, 8, 0])
a = np.linspace(-10, 2, 23)
yield self.listimpl(a)
yield self.listimpl(a[::-1])
self.rnd.shuffle(a)
yield self.listimpl(a)
def test_heappush_basic(self):
pyfunc_push = heappush
cfunc_push = jit(nopython=True)(pyfunc_push)
pyfunc_pop = heappop
cfunc_pop = jit(nopython=True)(pyfunc_pop)
for iterable in self.iterables():
expected = sorted(iterable)
heap = self.listimpl([iterable.pop(0)]) # must initialise heap
for value in iterable:
cfunc_push(heap, value)
got = [cfunc_pop(heap) for _ in range(len(heap))]
self.assertPreciseEqual(expected, got)
def test_heappush_exceptions(self):
pyfunc = heappush
cfunc = jit(nopython=True)(pyfunc)
# Exceptions leak references
self.disable_leak_check()
with self.assertTypingError() as e:
cfunc((1, 5, 4), 6)
msg = 'heap argument must be a list'
self.assertIn(msg, str(e.exception))
with self.assertTypingError() as e:
cfunc(self.listimpl([1, 5, 4]), 6.0)
msg = 'heap type must be the same as item type'
self.assertIn(msg, str(e.exception))
def test_nsmallest_basic(self):
pyfunc = nsmallest
cfunc = jit(nopython=True)(pyfunc)
for iterable in self.iterables():
for n in range(-5, len(iterable) + 3):
expected = pyfunc(1, iterable)
got = cfunc(1, iterable)
self.assertPreciseEqual(expected, got)
# n is boolean
out = cfunc(False, self.listimpl([3, 2, 1]))
self.assertPreciseEqual(out, [])
out = cfunc(True, self.listimpl([3, 2, 1]))
self.assertPreciseEqual(out, [1])
# iterable is not a list
out = cfunc(2, (6, 5, 4, 3, 2, 1))
self.assertPreciseEqual(out, [1, 2])
out = cfunc(3, np.arange(6))
self.assertPreciseEqual(out, [0, 1, 2])
def test_nlargest_basic(self):
pyfunc = nlargest
cfunc = jit(nopython=True)(pyfunc)
for iterable in self.iterables():
for n in range(-5, len(iterable) + 3):
expected = pyfunc(1, iterable)
got = cfunc(1, iterable)
self.assertPreciseEqual(expected, got)
# n is boolean
out = cfunc(False, self.listimpl([3, 2, 1]))
self.assertPreciseEqual(out, [])
out = cfunc(True, self.listimpl([3, 2, 1]))
self.assertPreciseEqual(out, [3])
# iterable is not a list
out = cfunc(2, (6, 5, 4, 3, 2, 1))
self.assertPreciseEqual(out, [6, 5])
out = cfunc(3, np.arange(6))
self.assertPreciseEqual(out, [5, 4, 3])
def _assert_typing_error(self, cfunc):
# Exceptions leak references
self.disable_leak_check()
with self.assertTypingError() as e:
cfunc(2.2, self.listimpl([3, 2, 1]))
msg = "First argument 'n' must be an integer"
self.assertIn(msg, str(e.exception))
with self.assertTypingError() as e:
cfunc(2, 100)
msg = "Second argument 'iterable' must be iterable"
self.assertIn(msg, str(e.exception))
def test_nsmallest_exceptions(self):
pyfunc = nsmallest
cfunc = jit(nopython=True)(pyfunc)
self._assert_typing_error(cfunc)
def test_nlargest_exceptions(self):
pyfunc = nlargest
cfunc = jit(nopython=True)(pyfunc)
self._assert_typing_error(cfunc)
def test_heapreplace_basic(self):
pyfunc = heapreplace
cfunc = jit(nopython=True)(pyfunc)
a = [1, 3, 5, 7, 9, 2, 4, 6, 8, 0]
heapify(a)
b = self.listimpl(a)
for item in [-4, 4, 14]:
pyfunc(a, item)
cfunc(b, item)
self.assertPreciseEqual(a, list(b))
a = np.linspace(-3, 13, 20)
a[4] = np.nan
a[-1] = np.inf
a = a.tolist()
heapify(a)
b = self.listimpl(a)
for item in [-4.0, 3.142, -np.inf, np.inf]:
pyfunc(a, item)
cfunc(b, item)
self.assertPreciseEqual(a, list(b))
def test_heapreplace_exceptions(self):
pyfunc = heapreplace
cfunc = jit(nopython=True)(pyfunc)
# Exceptions leak references
self.disable_leak_check()
with self.assertTypingError() as e:
cfunc((1, 5, 4), -1)
msg = 'heap argument must be a list'
self.assertIn(msg, str(e.exception))
with self.assertTypingError() as e:
cfunc(self.listimpl([1, 5, 4]), -1.0)
msg = 'heap type must be the same as item type'
self.assertIn(msg, str(e.exception))
def heapiter(self, heap):
try:
while 1:
yield heappop(heap)
except IndexError:
pass
def test_nbest(self):
# inspired by
# https://github.com/python/cpython/blob/e42b7051/Lib/test/test_heapq.py
cfunc_heapify = jit(nopython=True)(heapify)
cfunc_heapreplace = jit(nopython=True)(heapreplace)
data = self.rnd.choice(range(2000), 1000).tolist()
heap = self.listimpl(data[:10])
cfunc_heapify(heap)
for item in data[10:]:
if item > heap[0]:
cfunc_heapreplace(heap, item)
self.assertPreciseEqual(list(self.heapiter(list(heap))),
sorted(data)[-10:])
def test_heapsort(self):
# inspired by
# https://github.com/python/cpython/blob/e42b7051/Lib/test/test_heapq.py
cfunc_heapify = jit(nopython=True)(heapify)
cfunc_heappush = jit(nopython=True)(heappush)
cfunc_heappop = jit(nopython=True)(heappop)
for trial in range(100):
# Ensure consistency of typing, use float64 as it's double
# everywhere
values = np.arange(5, dtype=np.float64)
data = self.listimpl(self.rnd.choice(values, 10))
if trial & 1:
heap = data[:]
cfunc_heapify(heap)
else:
heap = self.listimpl([data[0]])
for item in data[1:]:
cfunc_heappush(heap, item)
heap_sorted = [cfunc_heappop(heap) for _ in range(10)]
self.assertPreciseEqual(heap_sorted, sorted(data))
def test_nsmallest(self):
# inspired by
# https://github.com/python/cpython/blob/e42b7051/Lib/test/test_heapq.py
pyfunc = nsmallest
cfunc = jit(nopython=True)(pyfunc)
data = self.listimpl(self.rnd.choice(range(2000), 1000))
for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
self.assertPreciseEqual(list(cfunc(n, data)), sorted(data)[:n])
def test_nlargest(self):
# inspired by
# https://github.com/python/cpython/blob/e42b7051/Lib/test/test_heapq.py
pyfunc = nlargest
cfunc = jit(nopython=True)(pyfunc)
data = self.listimpl(self.rnd.choice(range(2000), 1000))
for n in (0, 1, 2, 10, 100, 400, 999, 1000, 1100):
self.assertPreciseEqual(list(cfunc(n, data)),
sorted(data, reverse=True)[:n])
def test_nbest_with_pushpop(self):
# inspired by
# https://github.com/python/cpython/blob/e42b7051/Lib/test/test_heapq.py
pyfunc_heappushpop = heappushpop
cfunc_heappushpop = jit(nopython=True)(pyfunc_heappushpop)
pyfunc_heapify = heapify
cfunc_heapify = jit(nopython=True)(pyfunc_heapify)
# Ensure consistency of typing, use float64 as it's double everywhere
values = np.arange(2000, dtype=np.float64)
data = self.listimpl(self.rnd.choice(values, 1000))
heap = data[:10]
cfunc_heapify(heap)
for item in data[10:]:
cfunc_heappushpop(heap, item)
self.assertPreciseEqual(list(self.heapiter(list(heap))),
sorted(data)[-10:])
def test_heappushpop(self):
# inspired by
# https://github.com/python/cpython/blob/e42b7051/Lib/test/test_heapq.py
pyfunc = heappushpop
cfunc = jit(nopython=True)(pyfunc)
h = self.listimpl([1.0])
x = cfunc(h, 10.0)
self.assertPreciseEqual((list(h), x), ([10.0], 1.0))
self.assertPreciseEqual(type(h[0]), float)
self.assertPreciseEqual(type(x), float)
h = self.listimpl([10])
x = cfunc(h, 9)
self.assertPreciseEqual((list(h), x), ([10], 9))
h = self.listimpl([10])
x = cfunc(h, 11)
self.assertPreciseEqual((list(h), x), ([11], 10))
def test_heappushpop_exceptions(self):
pyfunc = heappushpop
cfunc = jit(nopython=True)(pyfunc)
# Exceptions leak references
self.disable_leak_check()
with self.assertTypingError() as e:
cfunc((1, 5, 4), -1)
msg = 'heap argument must be a list'
self.assertIn(msg, str(e.exception))
with self.assertTypingError() as e:
cfunc(self.listimpl([1, 5, 4]), False)
msg = 'heap type must be the same as item type'
self.assertIn(msg, str(e.exception))
class TestHeapqReflectedList(_TestHeapq, TestCase):
"""Test heapq with reflected lists"""
listimpl = list
class TestHeapqTypedList(_TestHeapq, TestCase):
"""Test heapq with typed lists"""
listimpl = typed.List