485 lines
14 KiB
Python
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
|