test_datamodel.py

from llvmlite import ir, binding as ll

from numba.core import types, datamodel
from numba.core.datamodel.testing import test_factory
from numba.core.datamodel.manager import DataModelManager
from numba.core.datamodel.models import OpaqueModel
import unittest


class TestBool(test_factory()):
    fe_type = types.boolean


class TestPyObject(test_factory()):
    fe_type = types.pyobject


class TestInt8(test_factory()):
    fe_type = types.int8


class TestInt16(test_factory()):
    fe_type = types.int16


class TestInt32(test_factory()):
    fe_type = types.int32


class TestInt64(test_factory()):
    fe_type = types.int64


class TestUInt8(test_factory()):
    fe_type = types.uint8


class TestUInt16(test_factory()):
    fe_type = types.uint16


class TestUInt32(test_factory()):
    fe_type = types.uint32


class TestUInt64(test_factory()):
    fe_type = types.uint64


class TestFloat(test_factory()):
    fe_type = types.float32


class TestDouble(test_factory()):
    fe_type = types.float64


class TestComplex(test_factory()):
    fe_type = types.complex64


class TestDoubleComplex(test_factory()):
    fe_type = types.complex128


class TestPointerOfInt32(test_factory()):
    fe_type = types.CPointer(types.int32)


class TestUniTupleOf2xInt32(test_factory()):
    fe_type = types.UniTuple(types.int32, 2)


class TestUniTupleEmpty(test_factory()):
    fe_type = types.UniTuple(types.int32, 0)


class TestTupleInt32Float32(test_factory()):
    fe_type = types.Tuple([types.int32, types.float32])


class TestTupleEmpty(test_factory()):
    fe_type = types.Tuple([])


class Test1DArrayOfInt32(test_factory()):
    fe_type = types.Array(types.int32, 1, 'C')


class Test2DArrayOfComplex128(test_factory()):
    fe_type = types.Array(types.complex128, 2, 'C')


class Test0DArrayOfInt32(test_factory()):
    fe_type = types.Array(types.int32, 0, 'C')


class TestArgInfo(unittest.TestCase):

    def _test_as_arguments(self, fe_args):
        """
        Test round-tripping types *fe_args* through the default data model's
        argument conversion and unpacking logic.
        """
        dmm = datamodel.default_manager
        fi = datamodel.ArgPacker(dmm, fe_args)

        module = ir.Module()
        fnty = ir.FunctionType(ir.VoidType(), [])
        function = ir.Function(module, fnty, name="test_arguments")
        builder = ir.IRBuilder()
        builder.position_at_end(function.append_basic_block())

        args = [ir.Constant(dmm.lookup(t).get_value_type(), None)
                for t in fe_args]

        # Roundtrip
        values = fi.as_arguments(builder, args)
        asargs = fi.from_arguments(builder, values)

        self.assertEqual(len(asargs), len(fe_args))
        valtys = tuple([v.type for v in values])
        self.assertEqual(valtys, fi.argument_types)

        expect_types = [a.type for a in args]
        got_types = [a.type for a in asargs]

        self.assertEqual(expect_types, got_types)

        # Assign names (check this doesn't raise)
        fi.assign_names(values, ["arg%i" for i in range(len(fe_args))])

        builder.ret_void()

        ll.parse_assembly(str(module))

    def test_int32_array_complex(self):
        fe_args = [types.int32,
                   types.Array(types.int32, 1, 'C'),
                   types.complex64]
        self._test_as_arguments(fe_args)

    def test_two_arrays(self):
        fe_args = [types.Array(types.int32, 1, 'C')] * 2
        self._test_as_arguments(fe_args)

    def test_two_0d_arrays(self):
        fe_args = [types.Array(types.int32, 0, 'C')] * 2
        self._test_as_arguments(fe_args)

    def test_tuples(self):
        fe_args = [types.UniTuple(types.int32, 2),
                   types.UniTuple(types.int32, 3)]
        self._test_as_arguments(fe_args)
        # Tuple of struct-likes
        arrty = types.Array(types.int32, 1, 'C')
        fe_args = [types.UniTuple(arrty, 2),
                   types.UniTuple(arrty, 3)]
        self._test_as_arguments(fe_args)
        # Nested tuple
        fe_args = [types.UniTuple(types.UniTuple(types.int32, 2), 3)]
        self._test_as_arguments(fe_args)

    def test_empty_tuples(self):
        # Empty tuple
        fe_args = [types.UniTuple(types.int16, 0),
                   types.Tuple(()),
                   types.int32]
        self._test_as_arguments(fe_args)

    def test_nested_empty_tuples(self):
        fe_args = [types.int32,
                   types.UniTuple(types.Tuple(()), 2),
                   types.int64]
        self._test_as_arguments(fe_args)


class TestMemInfo(unittest.TestCase):
    def setUp(self):
        self.dmm = datamodel.default_manager

    def test_number(self):
        ty = types.int32
        dm = self.dmm[ty]
        self.assertFalse(dm.contains_nrt_meminfo())

    def test_array(self):
        ty = types.int32[:]
        dm = self.dmm[ty]
        self.assertTrue(dm.contains_nrt_meminfo())

    def test_tuple_of_number(self):
        ty = types.UniTuple(dtype=types.int32, count=2)
        dm = self.dmm[ty]
        self.assertFalse(dm.contains_nrt_meminfo())

    def test_tuple_of_array(self):
        ty = types.UniTuple(dtype=types.int32[:], count=2)
        dm = self.dmm[ty]
        self.assertTrue(dm.contains_nrt_meminfo())


class TestMisc(unittest.TestCase):

    def test_issue2921(self):
        import numpy as np
        from numba import njit

        @njit
        def copy(a, b):
            for i in range(a.shape[0]):
                a[i] = b[i]

        b = np.arange(5, dtype=np.uint8).view(np.bool_)
        a = np.zeros_like(b)
        copy(a, b)
        np.testing.assert_equal(a, np.array((False,) + (True,) * 4))


class TestDMMChaining(unittest.TestCase):
    def test_basic(self):
        dmm = DataModelManager()

        class int_handler(OpaqueModel):
            pass

        class float_handler(OpaqueModel):
            pass

        dmm.register(types.Integer, int_handler)
        dmm.register(types.Float, float_handler)

        inter_dmm = DataModelManager()

        class new_int_handler(OpaqueModel):
            pass

        inter_dmm.register(types.Integer, new_int_handler)
        chained_dmm = inter_dmm.chain(dmm)

        # Check that the chained DMM has the new handler
        self.assertIsInstance(chained_dmm.lookup(types.intp), new_int_handler)
        # and not the old handler
        self.assertNotIsInstance(chained_dmm.lookup(types.intp), int_handler)
        # Check that the base DMM has the old handler
        self.assertIsInstance(dmm.lookup(types.intp), int_handler)
        # Check that float goes to the float_handler
        self.assertIsInstance(chained_dmm.lookup(types.float32), float_handler)
        self.assertIsInstance(dmm.lookup(types.float32), float_handler)
        # Check the intermediate DMM
        self.assertIsInstance(inter_dmm.lookup(types.intp), new_int_handler)
        with self.assertRaises(KeyError):
            inter_dmm.lookup(types.float32)


if __name__ == '__main__':
    unittest.main()