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

290 lines
8.6 KiB
Python
Raw Normal View History

2024-05-03 04:18:51 +03:00
import array
import numpy as np
from numba import jit
from numba.tests.support import TestCase, compile_function, MemoryLeakMixin
import unittest
@jit(nopython=True)
def len_usecase(buf):
return len(buf)
@jit(nopython=True)
def getitem_usecase(buf, i):
return buf[i]
@jit(nopython=True)
def getslice_usecase(buf, i, j):
s = buf[i:j]
return s[0] + 2 * s[-1]
@jit(nopython=True)
def setitem_usecase(buf, i, v):
buf[i] = v
@jit(nopython=True)
def iter_usecase(buf):
res = 0.0
for i, x in enumerate(buf):
res += x
res *= i + 1
return res
def attrgetter(attr):
code = """def func(x):
return x.%(attr)s
""" % locals()
pyfunc = compile_function("func", code, globals())
return jit(nopython=True)(pyfunc)
contiguous_usecase = attrgetter("contiguous")
c_contiguous_usecase = attrgetter("c_contiguous")
f_contiguous_usecase = attrgetter("f_contiguous")
itemsize_usecase = attrgetter("itemsize")
nbytes_usecase = attrgetter("nbytes")
ndim_usecase = attrgetter("ndim")
readonly_usecase = attrgetter("readonly")
shape_usecase = attrgetter("shape")
strides_usecase = attrgetter("strides")
class TestBufferProtocol(MemoryLeakMixin, TestCase):
"""
Test operations on buffer-providing objects.
"""
def _arrays(self):
n = 10
for letter, offset in [
('b', -3),
('B', 0),
('h', -5000),
('H', 40000),
('i', -100000),
('I', 1000000),
('l', -100000),
('L', 1000000),
('q', -2**60),
('Q', 2**63 + 1),
('f', 1.5),
('d', -1.5),
]:
yield array.array(letter, [i + offset for i in range(n)])
def _memoryviews(self):
n = 10
yield memoryview(bytearray(b"abcdefghi"))
yield memoryview(b"abcdefghi")
# Different item types
for dtype, start, stop in [
('int8', -10, 10),
('uint8', 0, 10),
('int16', -5000, 1000),
('uint16', 40000, 50000),
('int32', -100000, 100000),
('uint32', 0, 1000000),
('int64', -2**60, 10),
('uint64', 0, 2**64 - 10),
('float32', 1.5, 3.5),
('float64', 1.5, 3.5),
('complex64', -8j, 12 + 5j),
('complex128', -8j, 12 + 5j),
]:
yield memoryview(np.linspace(start, stop, n).astype(dtype))
# Different layouts
arr = np.arange(12).reshape((3, 4))
assert arr.flags.c_contiguous and not arr.flags.f_contiguous
yield memoryview(arr)
arr = arr.T
assert arr.flags.f_contiguous and not arr.flags.c_contiguous
yield memoryview(arr)
arr = arr[::2]
assert not arr.flags.f_contiguous and not arr.flags.c_contiguous
yield memoryview(arr)
def _readonlies(self):
yield b"xyz"
yield memoryview(b"abcdefghi")
arr = np.arange(5)
arr.setflags(write=False)
yield memoryview(arr)
def _check_unary(self, jitfunc, *args):
pyfunc = jitfunc.py_func
self.assertPreciseEqual(jitfunc(*args), pyfunc(*args))
def check_len(self, obj):
self._check_unary(len_usecase, obj)
def check_iter(self, obj):
self._check_unary(iter_usecase, obj)
def check_getitem(self, obj):
# Be careful to index all dimensions, since we don't support
# partial indexing yet.
def yield_indices(obj):
try:
shape = obj.shape
except AttributeError:
shape = len(obj),
for tup in np.ndindex(shape):
# Simple 1d buffer-providing objects usually don't support
# tuple indexing.
if len(tup) == 1:
yield tup[0]
else:
yield tup
for i in yield_indices(obj):
try:
expected = obj[i]
except (NotImplementedError, TypeError):
if isinstance(obj, memoryview):
# The memoryview object doesn't support all codes yet,
# fall back on the underlying object.
expected = obj.obj[i]
else:
raise
self.assertPreciseEqual(getitem_usecase(obj, i), expected)
def check_setitem(self, obj):
for i in range(len(obj)):
orig = list(obj)
val = obj[i] // 2 + 1
setitem_usecase(obj, i, val)
self.assertEqual(obj[i], val)
for j, val in enumerate(orig):
if j != i:
self.assertEqual(obj[j], val)
def check_getslice(self, obj):
self._check_unary(getslice_usecase, obj, 1, len(obj) - 1)
def test_len(self):
self.check_len(bytearray(5))
self.check_len(b"xyz")
for mem in self._memoryviews():
self.check_len(mem)
for arr in self._arrays():
self.check_len(arr)
for buf in self._readonlies():
self.check_getitem(buf)
def test_getitem(self):
self.check_getitem(bytearray(b"abc"))
self.check_getitem(b"xyz")
for mem in self._memoryviews():
self.check_getitem(mem)
for arr in self._arrays():
self.check_getitem(arr)
for buf in self._readonlies():
self.check_getitem(buf)
def test_getslice(self):
with self.assertTypingError():
self.check_getslice(bytearray(b"abcde"))
self.check_getslice(b"xyzuvw")
self.check_getslice(memoryview(b"xyzuvw"))
with self.assertTypingError():
self.check_getslice(array.array('i', range(10)))
for buf in self._readonlies():
self.check_getitem(buf)
def test_setitem(self):
self.check_setitem(bytearray(b"abcdefghi"))
for arr in self._arrays():
self.check_setitem(arr)
for mem in self._memoryviews():
self.check_getitem(mem)
# Read-only buffers
for buf in self._readonlies():
with self.assertTypingError():
self.check_setitem(buf)
def test_iter(self):
self.check_iter(bytearray(b"abc"))
self.check_iter(b"xyz")
self.check_iter(memoryview(b"xyz"))
for arr in self._arrays():
self.check_iter(arr)
for buf in self._readonlies():
self.check_getitem(buf)
class TestMemoryView(MemoryLeakMixin, TestCase):
"""
Test memoryview-specific attributes and operations.
"""
def _arrays(self):
arr = np.arange(12)
yield arr
arr = arr.reshape((3, 4))
yield arr
yield arr.T
yield arr[::2]
arr.setflags(write=False)
yield arr
arr = np.zeros(())
assert arr.ndim == 0
yield arr
def test_ndim(self):
for arr in self._arrays():
m = memoryview(arr)
self.assertPreciseEqual(ndim_usecase(m), arr.ndim)
def test_shape(self):
for arr in self._arrays():
m = memoryview(arr)
self.assertPreciseEqual(shape_usecase(m), arr.shape)
def test_strides(self):
for arr in self._arrays():
m = memoryview(arr)
self.assertPreciseEqual(strides_usecase(m), arr.strides)
def test_itemsize(self):
for arr in self._arrays():
m = memoryview(arr)
self.assertPreciseEqual(itemsize_usecase(m), arr.itemsize)
def test_nbytes(self):
for arr in self._arrays():
m = memoryview(arr)
self.assertPreciseEqual(nbytes_usecase(m), arr.size * arr.itemsize)
def test_readonly(self):
for arr in self._arrays():
m = memoryview(arr)
self.assertIs(readonly_usecase(m), not arr.flags.writeable)
m = memoryview(b"xyz")
self.assertIs(readonly_usecase(m), True)
m = memoryview(bytearray(b"xyz"))
self.assertIs(readonly_usecase(m), False)
def test_contiguous(self):
m = memoryview(bytearray(b"xyz"))
self.assertIs(contiguous_usecase(m), True)
self.assertIs(c_contiguous_usecase(m), True)
self.assertIs(f_contiguous_usecase(m), True)
for arr in self._arrays():
m = memoryview(arr)
# Note `arr.flags.contiguous` is wrong (it mimics c_contiguous)
self.assertIs(contiguous_usecase(m),
arr.flags.f_contiguous or arr.flags.c_contiguous)
self.assertIs(c_contiguous_usecase(m), arr.flags.c_contiguous)
self.assertIs(f_contiguous_usecase(m), arr.flags.f_contiguous)
if __name__ == '__main__':
unittest.main()