ai-content-maker/.venv/Lib/site-packages/numba/tests/test_optional.py

255 lines
6.3 KiB
Python
Raw Normal View History

2024-05-03 04:18:51 +03:00
import itertools
import numpy as np
import unittest
from numba import jit, njit
from numba.core import types
from numba.tests.support import TestCase
def return_double_or_none(x):
if x:
ret = None
else:
ret = 1.2
return ret
def return_different_statement(x):
if x:
return None
else:
return 1.2
def return_bool_optional_or_none(x, y):
if y:
z = False
else:
z = None
if x == 2:
# A boolean
return True
elif x == 1:
# A runtime optional
return z
else:
# None
return None
def is_this_a_none(x):
if x:
val_or_none = None
else:
val_or_none = x
if val_or_none is None:
return x - 1
if val_or_none is not None:
return x + 1
def a_is_b(a, b):
"""
Note in nopython mode, this operation does not make much sense.
Because we don't have objects anymore.
`a is b` is always False if not operating on None and Optional type
"""
return a is b
def a_is_not_b(a, b):
"""
This is `not (a is b)`
"""
return a is not b
class TestOptional(TestCase):
_numba_parallel_test_ = False
def test_return_double_or_none(self):
pyfunc = return_double_or_none
cfunc = njit((types.boolean,))(pyfunc)
for v in [True, False]:
self.assertPreciseEqual(pyfunc(v), cfunc(v))
def test_return_different_statement(self):
pyfunc = return_different_statement
cfunc = njit((types.boolean,))(pyfunc)
for v in [True, False]:
self.assertPreciseEqual(pyfunc(v), cfunc(v))
def test_return_bool_optional_or_none(self):
pyfunc = return_bool_optional_or_none
cfunc = njit((types.int32, types.int32,))(pyfunc)
for x, y in itertools.product((0, 1, 2), (0, 1)):
self.assertPreciseEqual(pyfunc(x, y), cfunc(x, y))
def test_is_this_a_none(self):
pyfunc = is_this_a_none
cfunc = njit((types.intp,))(pyfunc)
for v in [-1, 0, 1, 2]:
self.assertPreciseEqual(pyfunc(v), cfunc(v))
def test_is_this_a_none_objmode(self):
pyfunc = is_this_a_none
cfunc = jit((types.intp,), forceobj=True)(pyfunc)
self.assertTrue(cfunc.overloads[cfunc.signatures[0]].objectmode)
for v in [-1, 0, 1, 2]:
self.assertPreciseEqual(pyfunc(v), cfunc(v))
def test_a_is_b_intp(self):
pyfunc = a_is_b
cfunc = njit((types.intp, types.intp))(pyfunc)
# integer identity relies on `==`
self.assertTrue(cfunc(1, 1))
self.assertFalse(cfunc(1, 2))
def test_a_is_not_b_intp(self):
pyfunc = a_is_not_b
cfunc = njit((types.intp, types.intp))(pyfunc)
# integer identity relies on `==`
self.assertFalse(cfunc(1, 1))
self.assertTrue(cfunc(1, 2))
def test_optional_float(self):
def pyfunc(x, y):
if y is None:
return x
else:
return x + y
cfunc = njit("(float64, optional(float64))")(pyfunc)
self.assertAlmostEqual(pyfunc(1., 12.3), cfunc(1., 12.3))
self.assertAlmostEqual(pyfunc(1., None), cfunc(1., None))
def test_optional_array(self):
def pyfunc(x, y):
if y is None:
return x
else:
y[0] += x
return y[0]
cfunc = njit("(float32, optional(float32[:]))")(pyfunc)
cy = np.array([12.3], dtype=np.float32)
py = cy.copy()
self.assertAlmostEqual(pyfunc(1., py), cfunc(1., cy))
np.testing.assert_almost_equal(py, cy)
self.assertAlmostEqual(pyfunc(1., None), cfunc(1., None))
def test_optional_array_error(self):
def pyfunc(y):
return y[0]
cfunc = njit("(optional(int32[:]),)")(pyfunc)
with self.assertRaises(TypeError) as raised:
cfunc(None)
self.assertIn('expected array(int32, 1d, A), got None',
str(raised.exception))
y = np.array([0xabcd], dtype=np.int32)
self.assertEqual(cfunc(y), pyfunc(y))
def test_optional_array_attribute(self):
"""
Check that we can access attribute of an optional
"""
def pyfunc(arr, do_it):
opt = None
if do_it: # forces `opt` to be an optional of arr
opt = arr
return opt.shape[0]
cfunc = njit(pyfunc)
arr = np.arange(5)
self.assertEqual(pyfunc(arr, True), cfunc(arr, True))
def test_assign_to_optional(self):
"""
Check that we can assign to a variable of optional type
"""
@njit
def make_optional(val, get_none):
if get_none:
ret = None
else:
ret = val
return ret
@njit
def foo(val, run_second):
a = make_optional(val, True)
if run_second:
a = make_optional(val, False)
return a
self.assertIsNone(foo(123, False))
self.assertEqual(foo(231, True), 231)
def test_optional_thru_omitted_arg(self):
"""
Issue 1868
"""
def pyfunc(x=None):
if x is None:
x = 1
return x
cfunc = njit(pyfunc)
self.assertEqual(pyfunc(), cfunc())
self.assertEqual(pyfunc(3), cfunc(3))
def test_optional_unpack(self):
"""
Issue 2171
"""
def pyfunc(x):
if x is None:
return
else:
a, b = x
return a, b
tup = types.Tuple([types.intp] * 2)
opt_tup = types.Optional(tup)
sig = (opt_tup,)
cfunc = njit(sig)(pyfunc)
self.assertEqual(pyfunc(None), cfunc(None))
self.assertEqual(pyfunc((1, 2)), cfunc((1, 2)))
def test_many_optional_none_returns(self):
"""
Issue #4058
"""
@njit
def foo(maybe):
lx = None
if maybe:
lx = 10
return 1, lx
def work():
tmp = []
for _ in range(20000):
maybe = False
_ = foo(maybe)
# this caused "Fatal Python error: deallocating None" as there was no
# incref being made on the returned None.
work()
if __name__ == '__main__':
unittest.main()