from __future__ import print_function, absolute_import, division
import numpy
from numba import unittest_support as unittest
from numba.npyufunc.ufuncbuilder import UFuncBuilder, GUFuncBuilder
from numba import vectorize, guvectorize


def add(a, b):
    return a + b

def equals(a, b):
    return a == b

def guadd(a, b, c):
    x, y = c.shape
    for i in range(x):
        for j in range(y):
            c[i, j] = a[i, j] + b[i, j]


class TestUfuncBuilding(unittest.TestCase):
    def test_basic_ufunc(self):
        ufb = UFuncBuilder(add)
        ufb.add("int32(int32, int32)")
        ufb.add("int64(int64, int64)")
        ufunc = ufb.build_ufunc()

        a = numpy.arange(10, dtype='int32')
        b = ufunc(a, a)
        self.assertTrue(numpy.all(a + a == b))


    def test_ufunc_struct(self):
        ufb = UFuncBuilder(add)
        ufb.add("complex64(complex64, complex64)")
        ufunc = ufb.build_ufunc()

        a = numpy.arange(10, dtype='complex64') + 1j
        b = ufunc(a, a)
        self.assertTrue(numpy.all(a + a == b))
        self.assertTrue(b.dtype == numpy.dtype('complex64'))


class TestGUfuncBuilding(unittest.TestCase):
    def test_basic_gufunc(self):
        gufb = GUFuncBuilder(guadd, "(x, y),(x, y)->(x, y)")
        gufb.add("void(int32[:,:], int32[:,:], int32[:,:])")
        ufunc = gufb.build_ufunc()

        a = numpy.arange(10, dtype="int32").reshape(2, 5)
        b = ufunc(a, a)

        self.assertTrue(numpy.all(a + a == b))
        self.assertTrue(b.dtype == numpy.dtype('int32'))


    def test_gufunc_struct(self):
        gufb = GUFuncBuilder(guadd, "(x, y),(x, y)->(x, y)")
        gufb.add("void(complex64[:,:], complex64[:,:], complex64[:,:])")
        ufunc = gufb.build_ufunc()

        a = numpy.arange(10, dtype="complex64").reshape(2, 5) + 1j
        b = ufunc(a, a)

        self.assertTrue(numpy.all(a + a == b))
        self.assertTrue(b.dtype == numpy.dtype('complex64'))


class TestVectorizeDecor(unittest.TestCase):
    def test_vectorize(self):
        ufunc = vectorize(['int32(int32, int32)'])(add)
        a = numpy.arange(10, dtype='int32')
        b = ufunc(a, a)
        self.assertTrue(numpy.all(a + a == b))
        self.assertTrue(b.dtype == numpy.dtype('int32'))

    def test_vectorize_bool_return(self):
        ufunc = vectorize(['bool_(int32, int32)'])(equals)
        a = numpy.arange(10, dtype='int32')
        r = ufunc(a,a)
        self.assertTrue(numpy.all(r))
        self.assertTrue(r.dtype == numpy.dtype('bool_'))

    def test_guvectorize(self):
        ufunc = guvectorize(['(int32[:,:], int32[:,:], int32[:,:])'],
                            "(x,y),(x,y)->(x,y)")(guadd)
        a = numpy.arange(10, dtype='int32').reshape(2, 5)
        b = ufunc(a, a)
        self.assertTrue(numpy.all(a + a == b))
        self.assertTrue(b.dtype == numpy.dtype('int32'))

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

@@ -1058,36 +1058,37 @@
  self.assert
-True
+Equal
 (b.dtype
  == numpy.dt
@@ -1067,35 +1067,33 @@
 ertEqual(b.dtype
- ==
+,
  numpy.dtype('co
@@ -1486,36 +1486,37 @@
  self.assert
-True
+Equal
 (b.dtype
  == numpy.dt
@@ -1495,35 +1495,33 @@
 ertEqual(b.dtype
- ==
+,
  numpy.dtype('in
@@ -1887,36 +1887,37 @@
  self.assert
-True
+Equal
 (b.dtype
  == numpy.dt
@@ -1896,35 +1896,33 @@
 ertEqual(b.dtype
- ==
+,
  numpy.dtype('co
@@ -2195,36 +2195,37 @@
  self.assert
-True
+Equal
 (b.dtype
  == numpy.dt
@@ -2204,35 +2204,33 @@
 ertEqual(b.dtype
- ==
+,
  numpy.dtype('in
@@ -2466,20 +2466,21 @@
 sert
-True
+Equal
 (r.dtype
  == 
@@ -2471,27 +2471,25 @@
 qual(r.dtype
- ==
+,
  numpy.dtype
@@ -2811,20 +2811,21 @@
 sert
-True
+Equal
 (b.dtype
  == 
@@ -2820,19 +2820,17 @@
 (b.dtype
- ==
+,
  numpy.d

Use assertEqual to compare resulting types to expected types