def define_constructor(py_class, struct_typeclass, fields):
"""Define the jit-code constructor for `struct_typeclass` using the
Python type `py_class` and the required `fields`.
Use this instead of `define_proxy()` if the user does not want boxing
logic defined.
"""
# Build source code for the constructor
params = ', '.join(fields)
indent = ' ' * 8
init_fields_buf = []
for k in fields:
init_fields_buf.append(f"st.{k} = {k}")
init_fields = f'\n{indent}'.join(init_fields_buf)
source = f"""
def ctor({params}):
struct_type = struct_typeclass(list(zip({list(fields)}, [{params}])))
def impl({params}):
st = new(struct_type)
{init_fields}
return st
return impl
"""
glbs = dict(struct_typeclass=struct_typeclass, new=new)
exec(source, glbs)
ctor = glbs['ctor']
# Make it an overload
overload(py_class)(ctor)
def test_overload_inline_always_with_literally_in_inlinee(self):
# See issue #5887
def foo_ovld(dtype):
if not isinstance(dtype, types.StringLiteral):
def foo_noop(dtype):
return literally(dtype)
return foo_noop
if dtype.literal_value == 'str':
def foo_as_str_impl(dtype):
return 10
return foo_as_str_impl
if dtype.literal_value in ('int64', 'float64'):
def foo_as_num_impl(dtype):
return 20
return foo_as_num_impl
# define foo for literal str 'str'
def foo(dtype):
return 10
overload(foo, inline='always')(foo_ovld)
def test_impl(dtype):
return foo(dtype)
# check literal dispatch on 'str'
dtype = 'str'
self.check(test_impl, dtype, inline_expect={'foo': True})
# redefine foo to be correct for literal str 'int64'
def foo(dtype):
return 20
overload(foo, inline='always')(foo_ovld)
# check literal dispatch on 'int64'
dtype = 'int64'
self.check(test_impl, dtype, inline_expect={'foo': True})
def overload_elementwise_unary_ufunc(ufunc, name=None, dtype=None):
"""
Wrapper for unary ufuncs that returns an array
"""
if name is None:
name = ufunc.__name__
globals()[name] = ufunc
def unary_elementwise_ufunc_impl(a):
nb_dtype = determine_dtype(a, dtype)
typ = determine_input_type(a)
if isinstance(a, ArrayPointer):
def impl(a):
sz = len(a)
x = Array(sz, nb_dtype)
for i in range(sz):
# Convert the value to type "typ"
cast = typ(a[i])
x[i] = nb_dtype(ufunc(cast))
return x
return impl
else:
def impl(a):
# Convert the value to type typ
cast = typ(a)
return nb_dtype(ufunc(cast))
return impl
decorate = extending.overload(ufunc)
def wrapper(overload_func):
return decorate(unary_elementwise_ufunc_impl)
return wrapper
def overload_binary_op(op, inplace=False):
def omnisci_operator_impl(a, b):
if isinstance(a, ArrayPointer) and isinstance(b, ArrayPointer):
nb_dtype = a.eltype
def impl(a, b):
# XXX: raise exception if len(a) != len(b)
sz = len(a)
if inplace:
x = a
else:
x = Array(sz, nb_dtype)
for i in range(sz):
x[i] = nb_dtype(op(a[i], b[i]))
return x
return impl
decorate = extending.overload(op)
def wrapper(overload_func):
return decorate(omnisci_operator_impl)
return wrapper
def overload_binary_cmp_op(op):
def omnisci_operator_impl(a, e):
if isinstance(a, ArrayPointer):
if isinstance(e, ArrayPointer):
def impl(a, e):
if len(a) != len(e):
return False
for i in range(len(a)):
if not op(a[i], e[i]):
return False
return True
elif isinstance(e, types.Number):
def impl(a, e):
sz = len(a)
x = Array(sz, 'int8')
for i in range(sz):
x[i] = typesystem.boolean8(op(a[i], e))
return x
return impl
decorate = extending.overload(op)
def wrapper(overload_func):
return decorate(omnisci_operator_impl)
return wrapper
def wrapper(overload_func):
overload_func.__doc__ = self.format_docstring(
overload_func, ufunc_name, api)
functools.update_wrapper(fn, overload_func)
decorate = extending.overload(fn)
return decorate(unary_ufunc_impl)
def implements(self, func_name, api=API.ARRAY_API):
fn = self.create_function(func_name)
decorate = extending.overload(fn)
def wrapper(overload_func):
overload_func.__doc__ = self.format_docstring(
overload_func, func_name, api)
functools.update_wrapper(fn, overload_func)
return decorate(overload_func)
return wrapper
def expose_and_overload(func):
name = func.__name__
s = f'def {name}(*args, **kwargs): pass'
exec(s, globals())
fn = globals()[name]
decorate = extending.overload(fn)
def wrapper(overload_func):
return decorate(overload_func)
return wrapper
def not_implemented(self, func_name, api=API.ARRAY_API):
fn = self.create_function(func_name)
decorate = extending.overload(fn)
def unimplemented(*args, **kwargs):
raise errors.NumbaNotImplementedError(
f'Function "{func_name}" is not supported.\n'
'Please, open a ticket on the RBC project '
'and report this error if you need support for '
'this function.')
def wraps(overload_func):
original_doc = self.format_docstring(overload_func, func_name, api)
overload_func.__doc__ = f"❌ Not implemented\n{original_doc}"
functools.update_wrapper(fn, overload_func)
return decorate(unimplemented)
return wraps
def overload_unary_op(op):
def omnisci_operator_impl(a):
if isinstance(a, ArrayPointer):
nb_dtype = a.eltype
def impl(a):
sz = len(a)
x = Array(sz, nb_dtype)
for i in range(sz):
x[i] = nb_dtype(op(a[i]))
return x
return impl
decorate = extending.overload(op)
def wrapper(overload_func):
return decorate(omnisci_operator_impl)
return wrapper
def impl(a, b):
return val
return impl
elif a_none ^ b_none:
def impl(a, b):
return not val
return impl
return none_equality
overload(operator.eq)(gen_non_eq(True))
overload(operator.ne)(gen_non_eq(False))
#-------------------------------------------------------------------------------
@lower_getattr_generic(types.DeferredType)
def deferred_getattr(context, builder, typ, value, attr):
"""
Deferred.__getattr__ => redirect to the actual type.
"""
inner_type = typ.get()
val = context.cast(builder, value, typ, inner_type)
imp = context.get_getattr(inner_type, attr)
return imp(context, builder, inner_type, val, attr)
def overload_elementwise_binary_ufunc(ufunc, name=None, dtype=None):
"""
Wrapper for binary ufuncs that returns an array
"""
if name is None:
name = ufunc.__name__
globals()[name] = ufunc
def binary_ufunc_impl(a, b):
typA = determine_input_type(a)
typB = determine_input_type(b)
# XXX: raise error if len(a) != len(b)
@extending.register_jitable(_nrt=False)
def binary_impl(a, b, nb_dtype):
sz = len(a)
x = Array(sz, nb_dtype)
for i in range(sz):
cast_a = typA(a[i])
cast_b = typB(b[i])
x[i] = nb_dtype(ufunc(cast_a, cast_b))
return x
@extending.register_jitable(_nrt=False)
def broadcast(e, sz, dtype):
b = Array(sz, dtype)
b.fill(e)
return b
if isinstance(a, ArrayPointer) and isinstance(b, ArrayPointer):
nb_dtype = determine_dtype(a, dtype)
def impl(a, b):
return binary_impl(a, b, nb_dtype)
return impl
elif isinstance(a, ArrayPointer):
nb_dtype = determine_dtype(a, dtype)
other_dtype = b
def impl(a, b):
b = broadcast(b, len(a), other_dtype)
return binary_impl(a, b, nb_dtype)
return impl
elif isinstance(b, ArrayPointer):
nb_dtype = determine_dtype(b, dtype)
other_dtype = a
def impl(a, b):
a = broadcast(a, len(b), other_dtype)
return binary_impl(a, b, nb_dtype)
return impl
else:
nb_dtype = determine_dtype(a, dtype)
def impl(a, b):
cast_a = typA(a)
cast_b = typB(b)
return nb_dtype(ufunc(cast_a, cast_b))
return impl
decorate = extending.overload(ufunc)
def wrapper(overload_func):
return decorate(binary_ufunc_impl)
return wrapper
