Viewing File: /home/ubuntu/combine_ai/combine/lib/python3.10/site-packages/numba/tests/test_unpack_sequence.py

import numpy as np

import unittest
from numba import jit, njit
from numba.core import errors, types
from numba import typeof
from numba.tests.support import TestCase, MemoryLeakMixin
from numba.tests.support import no_pyobj_flags as nullary_no_pyobj_flags
from numba.tests.support import force_pyobj_flags as nullary_force_pyobj_flags

force_pyobj_flags = {'forceobj': True}
no_pyobj_flags = {'nopython': True}


def unpack_list(l):
    a, b, c = l
    return (a, b, c)


def unpack_shape(a):
    x, y, z = a.shape
    return x + y + z


def unpack_range():
    a, b, c = range(3)
    return a + b + c


def unpack_range_too_small():
    a, b, c = range(2)
    return a + b + c


def unpack_range_too_large():
    a, b, c = range(4)
    return a + b + c


def unpack_tuple():
    a, b, c = (1, 2, 3)
    return a + b + c


def unpack_tuple_too_small():
    a, b, c = (1, 2)
    return a + b + c


def unpack_tuple_too_large():
    a, b, c = (1, 2, 3, 4)
    return a + b + c


def unpack_heterogeneous_tuple_too_small():
    a, b, c = (1, 2.5j)
    return a + b + c


def unpack_heterogeneous_tuple_too_large():
    a, b, c = (1, 2.5, 3j, 4)
    return a + b + c


def unpack_heterogeneous_tuple():
    a, b, c = (1, 2.5, 3j)
    return a + b + c


def unpack_nested_heterogeneous_tuple():
    a, (b, c) = (1, (2.5, 3j))
    return a + b + c


def unpack_arbitrary(seq):
    a, b = seq
    return b, a


def unpack_nrt():
    a = np.zeros(1)
    b = np.zeros(2)
    tup = b, a
    alpha, beta = tup
    return alpha, beta


def chained_unpack_assign1(x, y):
    # Used to fail in object mode (issue #580)
    a = (b, c) = (x, y)
    (d, e) = a
    return d + e + b + c


def conditional_swap(x, y):
    # Used to produce invalid code (issue #977)
    if x > 0:
        x, y = y, x
    return x, y


class TestUnpack(MemoryLeakMixin, TestCase):

    def check_nullary_npm(self, pyfunc):
        cfunc = njit(pyfunc)
        self.assertPreciseEqual(cfunc(), pyfunc())

    def check_nullary_objmode(self, pyfunc):
        cfunc = jit(forceobj=True)(pyfunc)
        self.assertPreciseEqual(cfunc(), pyfunc())

    def test_unpack_list(self):
        pyfunc = unpack_list
        cfunc = jit(forceobj=True)(pyfunc)
        l = [1, 2, 3]
        self.assertEqual(cfunc(l), pyfunc(l))

    def test_unpack_shape(self):
        pyfunc = unpack_shape
        cfunc = jit((types.Array(dtype=types.int32, ndim=3, layout='C'),),
                    forceobj=True)(pyfunc)
        a = np.zeros(shape=(1, 2, 3)).astype(np.int32)
        self.assertPreciseEqual(cfunc(a), pyfunc(a))

    def test_unpack_shape_npm(self):
        pyfunc = unpack_shape
        cfunc = njit((types.Array(dtype=types.int32, ndim=3, layout='C'),),
                     )(pyfunc)
        a = np.zeros(shape=(1, 2, 3)).astype(np.int32)
        self.assertPreciseEqual(cfunc(a), pyfunc(a))

    def test_unpack_range(self):
        self.check_nullary_objmode(unpack_range)

    def test_unpack_range_npm(self):
        self.check_nullary_npm(unpack_range)

    def test_unpack_tuple(self):
        self.check_nullary_objmode(unpack_tuple)

    def test_unpack_tuple_npm(self):
        self.check_nullary_npm(unpack_tuple)

    def test_unpack_heterogeneous_tuple(self):
        self.check_nullary_objmode(unpack_heterogeneous_tuple)

    def test_unpack_heterogeneous_tuple_npm(self):
        self.check_nullary_npm(unpack_heterogeneous_tuple)

    def test_unpack_nested_heterogeneous_tuple(self):
        self.check_nullary_objmode(unpack_nested_heterogeneous_tuple)

    def test_unpack_nested_heterogeneous_tuple_npm(self):
        self.check_nullary_npm(unpack_nested_heterogeneous_tuple)

    def test_chained_unpack_assign(self, flags=force_pyobj_flags):
        pyfunc = chained_unpack_assign1
        cfunc = jit((types.int32, types.int32), **flags)(pyfunc)
        args = (4, 5)
        self.assertPreciseEqual(cfunc(*args), pyfunc(*args))

    def test_chained_unpack_assign_npm(self):
        self.test_chained_unpack_assign(flags=no_pyobj_flags)

    def check_unpack_error(self, pyfunc, flags=force_pyobj_flags,
                           exc=ValueError):
        with self.assertRaises(exc):
            cfunc = jit((), **flags)(pyfunc)
            cfunc()

    def test_unpack_tuple_too_small(self):
        self.check_unpack_error(unpack_tuple_too_small)
        self.check_unpack_error(unpack_heterogeneous_tuple_too_small)

    def test_unpack_tuple_too_small_npm(self):
        self.check_unpack_error(unpack_tuple_too_small, no_pyobj_flags,
                                errors.TypingError)
        self.check_unpack_error(unpack_heterogeneous_tuple_too_small,
                                no_pyobj_flags, errors.TypingError)

    def test_unpack_tuple_too_large(self):
        self.check_unpack_error(unpack_tuple_too_large)
        self.check_unpack_error(unpack_heterogeneous_tuple_too_large)

    def test_unpack_tuple_too_large_npm(self):
        self.check_unpack_error(unpack_tuple_too_large, no_pyobj_flags,
                                errors.TypingError)
        self.check_unpack_error(unpack_heterogeneous_tuple_too_large,
                                no_pyobj_flags, errors.TypingError)

    def test_unpack_range_too_small(self):
        self.check_unpack_error(unpack_range_too_small)

    def test_unpack_range_too_small_npm(self):
        self.check_unpack_error(unpack_range_too_small, no_pyobj_flags)

    def test_unpack_range_too_large(self):
        self.check_unpack_error(unpack_range_too_large)

    def test_unpack_range_too_large_npm(self):
        self.check_unpack_error(unpack_range_too_large, no_pyobj_flags)

    def check_conditional_swap(self, flags=force_pyobj_flags):
        cfunc = jit((types.int32, types.int32), **flags)(conditional_swap)
        self.assertPreciseEqual(cfunc(4, 5), (5, 4))
        self.assertPreciseEqual(cfunc(0, 5), (0, 5))

    def test_conditional_swap(self):
        self.check_conditional_swap()

    def test_conditional_swap_npm(self):
        self.check_conditional_swap(no_pyobj_flags)

    def test_unpack_tuple_of_arrays(self):
        tup = tuple(np.zeros(i + 1) for i in range(2))
        tupty = typeof(tup)
        pyfunc = unpack_arbitrary
        cfunc = njit((tupty,))(pyfunc)
        self.assertPreciseEqual(cfunc(tup), pyfunc(tup))

    def test_unpack_nrt(self):
        pyfunc = unpack_nrt
        cfunc = njit((),)(pyfunc)
        self.assertPreciseEqual(cfunc(), pyfunc())

    def test_invalid_unpack(self):
        pyfunc = unpack_arbitrary
        with self.assertRaises(errors.TypingError) as raises:
            njit((types.int32,))(pyfunc)
        self.assertIn("failed to unpack int32", str(raises.exception))


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