test_slices.py

from functools import partial
import itertools
from itertools import chain, product, starmap
import sys

import numpy as np

from numba import jit, literally, njit, typeof, TypingError
from numba.core import utils, types
from numba.tests.support import TestCase, MemoryLeakMixin
from numba.core.types.functions import _header_lead
import unittest


def slice_passing(sl):
    return sl.start, sl.stop, sl.step

def slice_constructor(*args):
    sl = slice(*args)
    return sl.start, sl.stop, sl.step

def slice_construct_and_use(args, l):
    sl = slice(*args)
    return l[sl]

def slice_indices(s, *indargs):
    return s.indices(*indargs)

class TestSlices(MemoryLeakMixin, TestCase):

    def test_slice_passing(self):
        """
        Check passing a slice object to a Numba function.
        """
        # NOTE this also checks slice attributes
        def check(a, b, c, d, e, f):
            sl = slice(a, b, c)
            got = cfunc(sl)
            self.assertPreciseEqual(got, (d, e, f))

        maxposint = sys.maxsize
        maxnegint = -maxposint - 1
        cfunc = jit(nopython=True)(slice_passing)

        # Positive steps
        start_cases = [(None, 0), (42, 42), (-1, -1)]
        stop_cases = [(None, maxposint), (9, 9), (-11, -11)]
        step_cases = [(None, 1), (12, 12)]
        for (a, d), (b, e), (c, f) in itertools.product(start_cases,
                                                        stop_cases,
                                                        step_cases):
            check(a, b, c, d, e, f)

        # Negative steps
        start_cases = [(None, maxposint), (42, 42), (-1, -1)]
        stop_cases = [(None, maxnegint), (9, 9), (-11, -11)]
        step_cases = [(-1, -1), (-12, -12)]
        for (a, d), (b, e), (c, f) in itertools.product(start_cases,
                                                        stop_cases,
                                                        step_cases):
            check(a, b, c, d, e, f)

        # Some member is neither integer nor None
        with self.assertRaises(TypeError):
            cfunc(slice(1.5, 1, 1))

    def test_slice_constructor(self):
        """
        Test the 'happy path' for slice() constructor in nopython mode.
        """
        maxposint = sys.maxsize
        maxnegint = -maxposint - 1
        a = np.arange(10)
        cfunc = jit(nopython=True)(slice_constructor)
        cfunc_use = jit(nopython=True)(slice_construct_and_use)
        for args, expected in [
            ((None,), (0, maxposint, 1)),
            ((5,), (0, 5, 1)),
            ((None, None), (0, maxposint, 1)),
            ((1, None), (1, maxposint, 1)),
            ((None, 2), (0, 2, 1)),
            ((1, 2), (1, 2, 1)),
            ((None, None, 3), (0, maxposint, 3)),
            ((None, 2, 3), (0, 2, 3)),
            ((1, None, 3), (1, maxposint, 3)),
            ((1, 2, 3), (1, 2, 3)),
            ((None, None, -1), (maxposint, maxnegint, -1)),
            ((10, None, -1), (10, maxnegint, -1)),
            ((None, 5, -1), (maxposint, 5, -1)),
            ((10, 5, -1), (10, 5, -1)),
        ]:
            got = cfunc(*args)
            self.assertPreciseEqual(got, expected)
            usage = slice_construct_and_use(args, a)
            cusage = cfunc_use(args, a)
            self.assertPreciseEqual(usage, cusage)

    def test_slice_constructor_cases(self):
        """
        Test that slice constructor behaves same in python and compiled code.
        """
        options = (None, -1, 0, 1)
        arg_cases = chain.from_iterable(
            product(options, repeat=n) for n in range(5)
        )
        array = np.arange(10)

        cfunc = jit(nopython=True)(slice_construct_and_use)

        self.disable_leak_check()
        for args in arg_cases:
            try:
                expected = slice_construct_and_use(args, array)
            except TypeError as py_type_e:
                # Catch cases of 0, or more than 3 arguments.
                # This becomes a typing error in numba
                n_args = len(args)
                self.assertRegexpMatches(
                    str(py_type_e),
                    r"slice expected at (most|least) (3|1) arguments?, got {}"
                    .format(n_args)
                )
                with self.assertRaises(TypingError) as numba_e:
                    cfunc(args, array)
                self.assertIn(
                    _header_lead,
                    str(numba_e.exception)
                )
                self.assertIn(
                    ", ".join(str(typeof(arg)) for arg in args),
                    str(numba_e.exception)
                )
            except Exception as py_e:
                with self.assertRaises(type(py_e)) as numba_e:
                    cfunc(args, array)
                self.assertIn(
                    str(py_e),
                    str(numba_e.exception)
                )
            else:
                self.assertPreciseEqual(expected, cfunc(args, array))

    def test_slice_indices(self):
        """Test that a numba slice returns same result for .indices as a python one."""
        slices = starmap(
            slice,
            product(
                chain(range(-5, 5), (None,)),
                chain(range(-5, 5), (None,)),
                chain(range(-5, 5), (None,))
            )
        )
        lengths = range(-2, 3)

        cfunc = jit(nopython=True)(slice_indices)

        for s, l in product(slices, lengths):
            try:
                expected = slice_indices(s, l)
            except Exception as py_e:
                with self.assertRaises(type(py_e)) as numba_e:
                    cfunc(s, l)
                self.assertIn(
                    str(py_e),
                    str(numba_e.exception)
                )
            else:
                self.assertPreciseEqual(expected, cfunc(s, l))

    def test_slice_indices_examples(self):
        """Tests for specific error cases."""
        cslice_indices = jit(nopython=True)(slice_indices)

        with self.assertRaises(TypingError) as e:
            cslice_indices(slice(None), 1, 2, 3)
        self.assertIn(
             "indices() takes exactly one argument (3 given)",
             str(e.exception)
        )

        with self.assertRaises(TypingError) as e:
            cslice_indices(slice(None, None, 0), 1.2)
        self.assertIn(
            "'%s' object cannot be interpreted as an integer" % typeof(1.2),
            str(e.exception)
        )

    def test_slice_from_constant(self):
        test_tuple = (1, 2, 3, 4)

        for ts in itertools.product(
            [None, 1, 2, 3], [None, 1, 2, 3], [None, 1, 2, -1, -2]
        ):
            ts = slice(*ts)

            @jit(nopython=True)
            def test_fn():
                return test_tuple[ts]

            self.assertEqual(test_fn(), test_fn.py_func())

    def test_literal_slice_distinct(self):
        sl1 = types.misc.SliceLiteral(slice(1, None, None))
        sl2 = types.misc.SliceLiteral(slice(None, None, None))
        sl3 = types.misc.SliceLiteral(slice(1, None, None))

        self.assertNotEqual(sl1, sl2)
        self.assertEqual(sl1, sl3)

    def test_literal_slice_boxing(self):
        # Tests that a literal slice can be
        # returned from a JIT function.
        @njit
        def f(x):
            return literally(x)

        slices = (
            slice(1, 4, 2),
            slice(1, 2),
            slice(1),
            slice(None, 1, 1),
            slice(1, None, 1),
            slice(None, None, 1),
            slice(None),
            slice(None, None, None)
        )
        for sl in slices:
            self.assertEqual(sl, f(sl))


    def test_literal_slice_freevar(self):
        # Tests passing a literal slice as a freevar
        # in a closure.
        z = slice(1, 2, 3)
        @njit
        def foo():
            return z

        self.assertEqual(z, foo())

    def test_literal_slice_maxint(self):
        # Tests that passing a slice with an integer
        # that exceeds the maxint size throws a reasonable
        # error message.
        @njit()
        def foo(z):
            return literally(z)

        maxval = int(2**63)
        with self.assertRaises(ValueError) as e:
            foo(slice(None, None, -maxval-1))
        self.assertIn(
            "Int value is too large",
            str(e.exception)
        )


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