test_unsafe_intrinsics.py

import random
import numpy as np

from numba.tests.support import TestCase, captured_stdout
from numba import njit, literally
from numba.core import types
from numba.cpython.unsafe.tuple import tuple_setitem, build_full_slice_tuple
from numba.np.unsafe.ndarray import to_fixed_tuple, empty_inferred
from numba.core.unsafe.bytes import memcpy_region
from numba.core.unsafe.refcount import dump_refcount
from numba.cpython.unsafe.numbers import trailing_zeros, leading_zeros
from numba.core.errors import TypingError


class TestTupleIntrinsic(TestCase):
    """Tests for numba.unsafe.tuple
    """
    def test_tuple_setitem(self):
        @njit
        def foo(tup, idxs, vals):
            out_tup = tup
            for i, v in zip(idxs, vals):
                out_tup = tuple_setitem(out_tup, i, v)
            return tup, out_tup

        random.seed(123)
        for _ in range(20):
            # Random data
            n = random.randint(1, 10)
            tup = tuple([random.randint(0, n) for i in range(n)])
            vals = tuple([random.randint(10, 20) for i in range(n)])
            idxs = list(range(len(vals)))
            random.shuffle(idxs)
            idxs = tuple(idxs)
            # Expect
            expect_tup = tuple(tup)
            expect_out = np.asarray(expect_tup)
            expect_out[np.asarray(idxs)] = vals
            # Got
            got_tup, got_out = foo(tup, idxs, vals)
            # Check
            self.assertEqual(got_tup, expect_tup)
            self.assertEqual(got_out, tuple(expect_out))

    def test_slice_tuple(self):
        @njit
        def full_slice_array(a, n):
            # Since numba slices can't be boxed at the moment
            return a[build_full_slice_tuple(literally(n))]

        for n in range(1, 3):
            a = np.random.random(np.arange(n) + 1)
            for i in range(1, n + 1):
                np.testing.assert_array_equal(a, full_slice_array(a, i))
            with self.assertRaises(TypingError):
                # numpy would throw an IndexError here
                full_slice_array(a, n + 1)


class TestNdarrayIntrinsic(TestCase):
    """Tests for numba.unsafe.ndarray
    """
    def test_to_fixed_tuple(self):
        const = 3

        @njit
        def foo(array):
            a = to_fixed_tuple(array, length=1)
            b = to_fixed_tuple(array, 2)
            c = to_fixed_tuple(array, const)
            d = to_fixed_tuple(array, 0)
            return a, b, c, d

        np.random.seed(123)
        for _ in range(10):
            # Random data
            arr = np.random.random(3)
            # Run
            a, b, c, d = foo(arr)
            # Check
            self.assertEqual(a, tuple(arr[:1]))
            self.assertEqual(b, tuple(arr[:2]))
            self.assertEqual(c, tuple(arr[:3]))
            self.assertEqual(d, ())

        # Check error with ndim!=1
        with self.assertRaises(TypingError) as raises:
            foo(np.random.random((1, 2)))
        self.assertIn("Not supported on array.ndim=2",
                      str(raises.exception))

        # Check error with non-constant length
        @njit
        def tuple_with_length(array, length):
            return to_fixed_tuple(array, length)

        with self.assertRaises(TypingError) as raises:
            tuple_with_length(np.random.random(3), 1)
        expectmsg = "*length* argument must be a constant"
        self.assertIn(expectmsg, str(raises.exception))

    def test_issue_3586_variant1(self):
        @njit
        def func():
            S = empty_inferred((10,))
            a = 1.1
            for i in range(len(S)):
                S[i] = a + 2
            return S

        got = func()
        expect = np.asarray([3.1] * 10)
        np.testing.assert_array_equal(got, expect)

    def test_issue_3586_variant2(self):
        @njit
        def func():
            S = empty_inferred((10,))
            a = 1.1
            for i in range(S.size):
                S[i] = a + 2
            return S

        got = func()
        expect = np.asarray([3.1] * 10)
        np.testing.assert_array_equal(got, expect)


class TestBytesIntrinsic(TestCase):
    """Tests for numba.unsafe.bytes
    """
    def test_memcpy_region(self):
        @njit
        def foo(dst, dst_index, src, src_index, nbytes):
            # last arg is assume 1 byte alignment
            memcpy_region(dst.ctypes.data, dst_index,
                          src.ctypes.data, src_index, nbytes, 1)

        d = np.zeros(10, dtype=np.int8)
        s = np.arange(10, dtype=np.int8)

        # copy s[1:6] to d[4:9]
        foo(d, 4, s, 1, 5)

        expected = [0, 0, 0, 0, 1, 2, 3, 4, 5, 0]
        np.testing.assert_array_equal(d, expected)


class TestRefCount(TestCase):
    def test_dump_refcount(self):
        @njit
        def use_dump_refcount():
            a = np.ones(10)
            b = (a, a)
            dump_refcount(a)
            dump_refcount(b)

        # Capture output to sys.stdout
        with captured_stdout() as stream:
            use_dump_refcount()

        output = stream.getvalue()
        # Check that it printed
        pat = "dump refct of {}"
        aryty = types.float64[::1]
        tupty = types.Tuple.from_types([aryty] * 2)
        self.assertIn(pat.format(aryty), output)
        self.assertIn(pat.format(tupty), output)


class TestZeroCounts(TestCase):
    def test_zero_count(self):
        lz = njit(lambda x: leading_zeros(x))
        tz = njit(lambda x: trailing_zeros(x))

        evens = [2, 42, 126, 128]

        for T in types.unsigned_domain:
            self.assertTrue(tz(T(0)) == lz(T(0)) == T.bitwidth)
            for i in range(T.bitwidth):
                val = T(2 ** i)
                self.assertEqual(lz(val) + tz(val) + 1, T.bitwidth)
            for n in evens:
                self.assertGreater(tz(T(n)), 0)
                self.assertEqual(tz(T(n + 1)), 0)

        for T in types.signed_domain:
            self.assertTrue(tz(T(0)) == lz(T(0)) == T.bitwidth)
            for i in range(T.bitwidth - 1):
                val = T(2 ** i)
                self.assertEqual(lz(val) + tz(val) + 1, T.bitwidth)
                self.assertEqual(lz(-val), 0)
                self.assertEqual(tz(val), tz(-val))
            for n in evens:
                self.assertGreater(tz(T(n)), 0)
                self.assertEqual(tz(T(n + 1)), 0)

    def check_error_msg(self, func):
        cfunc = njit(lambda *x: func(*x))
        func_name = func._name

        unsupported_types = filter(
            lambda x: not isinstance(x, types.Integer), types.number_domain
        )
        for typ in sorted(unsupported_types, key=str):
            with self.assertRaises(TypingError) as e:
                cfunc(typ(2))
            self.assertIn(
                "{} is only defined for integers, but value passed was '{}'."
                .format(func_name, typ),
                str(e.exception),
            )

        # Testing w/ too many/few arguments
        def check(args, string):
            with self.assertRaises((TypingError, TypeError)) as e:
                cfunc(*args)
            self.assertIn(
                "{}() ".format(func_name),
                str(e.exception)
            )

        check((1, 2), "takes 2 positional arguments but 3 were given")
        check((), "missing 1 required positional argument")

    def test_trailing_zeros_error(self):
        self.check_error_msg(trailing_zeros)

    def test_leading_zeros_error(self):
        self.check_error_msg(leading_zeros)