import typing
from typing import Annotated, Any, get_type_hints, get_args
from numba import njit, guvectorize, vectorize, __version__
import numpy as np
import inspect
import functools
__version__
'0.54.1'
def parse_func(f):
try:
print('type_hints', get_type_hints(f, include_extras=True), sep='\n', end='\n\n')
except:
print('type_hints', '<not available>', sep='\n', end='\n\n')
try:
print('argspec', inspect.getfullargspec(f), sep='\n', end='\n\n')
except:
print('argspec', '<not available>', sep='\n', end='\n\n')
try:
print('docstring', f.__doc__, sep='\n', end='\n\n')
except:
print('docstring', '<not available>', sep='\n', end='\n\n')
x_type = Annotated[Any, "custom_x"]
y_type = Annotated[Any, "custom_y"]
z_type = Annotated[Any, "custom_z"]
def py_func(x: x_type, y: y_type) -> z_type:
"""py_func"""
return x+y
parse_func(py_func)
type_hints {'x': typing.Annotated[typing.Any, 'custom_x'], 'y': typing.Annotated[typing.Any, 'custom_y'], 'return': typing.Annotated[typing.Any, 'custom_z']} argspec FullArgSpec(args=['x', 'y'], varargs=None, varkw=None, defaults=None, kwonlyargs=[], kwonlydefaults=None, annotations={'return': typing.Annotated[typing.Any, 'custom_z'], 'x': typing.Annotated[typing.Any, 'custom_x'], 'y': typing.Annotated[typing.Any, 'custom_y']}) docstring py_func
@njit
def nb_njit_func(x: x_type, y: y_type) -> z_type:
"""nb_njit_func"""
return x+y
parse_func(nb_njit_func.py_func)
type_hints {'x': typing.Annotated[typing.Any, 'custom_x'], 'y': typing.Annotated[typing.Any, 'custom_y'], 'return': typing.Annotated[typing.Any, 'custom_z']} argspec FullArgSpec(args=['x', 'y'], varargs=None, varkw=None, defaults=None, kwonlyargs=[], kwonlydefaults=None, annotations={'return': typing.Annotated[typing.Any, 'custom_z'], 'x': typing.Annotated[typing.Any, 'custom_x'], 'y': typing.Annotated[typing.Any, 'custom_y']}) docstring nb_njit_func
parse_func(nb_njit_func)
type_hints <not available> argspec <not available> docstring nb_njit_func
@guvectorize
def nb_guvec_func(x: x_type, y: y_type) -> z_type:
"""nb_guvec_func"""
return x+y
parse_func(nb_guvec_func)
type_hints {} argspec FullArgSpec(args=['func'], varargs=None, varkw=None, defaults=None, kwonlyargs=[], kwonlydefaults=None, annotations={}) docstring None
@vectorize
def nb_vec_func(x: x_type, y: y_type) -> z_type:
"""nb_vec_func"""
return x+y
parse_func(nb_vec_func)
type_hints <not available> argspec <not available> docstring nb_vec_func
annotations = {}
def parse_annotations(f):
inputs = []
outputs = []
for k,v in get_type_hints(f, include_extras=True).items():
datatype, component = get_args(v)
if k == 'return':
outputs.append(component)
else:
inputs.append(component)
annotations[f.__name__] = (inputs, outputs)
return f
@vectorize
@parse_annotations
def nb_vec_func(x: x_type, y: y_type) -> z_type:
"""nb_vec_func"""
return x+y
annotations
{'nb_vec_func': (['custom_x', 'custom_y'], ['custom_z'])}
def wrap_vectorize(*args, **kwargs):
if len(args) == 1 and not kwargs and callable(args[0]):
func, = args
args = tuple()
has_args = False
else:
has_args = True
def outer(func):
wrapped_vec = vectorize(*args, **kwargs)(func)
functools.update_wrapper(wrapped_vec, func)
return wrapped_vec
if has_args:
return outer
else:
return outer(func)
@wrap_vectorize([
'float32(float32, uint8)',
'float64(float64, uint8)',
], fastmath=True)
def nb_vec_func_wrapped(x: x_type, y: y_type) -> z_type:
"""nb_vec_func"""
return x+y
parse_func(nb_vec_func_wrapped)
type_hints {'x': typing.Annotated[typing.Any, 'custom_x'], 'y': typing.Annotated[typing.Any, 'custom_y'], 'return': typing.Annotated[typing.Any, 'custom_z']} argspec <not available> docstring nb_vec_func
type(nb_vec_func_wrapped)
numba.np.ufunc.dufunc.DUFunc
nb_vec_func_wrapped.__wrapped__
<function __main__.nb_vec_func_wrapped(x: Annotated[Any, 'custom_x'], y: Annotated[Any, 'custom_y']) -> Annotated[Any, 'custom_z']>
# call function with correct datatypes [float, int]
nb_vec_func_wrapped(np.random.randn(2,2), 2)
array([[-0.14655443, 1.84584324], [ 1.62229537, 1.32742006]])
# call function with incorrect datatypes [float, float]
# should raise TypeError
nb_vec_func_wrapped(np.random.randn(2,2), 2.5)
--------------------------------------------------------------------------- TypeError Traceback (most recent call last) ~\AppData\Local\Temp/ipykernel_2716/3192411579.py in <module> 1 # call function with incorrect datatypes [float, float], should raise TypeError ----> 2 nb_vec_func_wrapped(np.random.randn(2,2), 2.5) TypeError: ufunc 'nb_vec_func_wrapped' not supported for the input types, and the inputs could not be safely coerced to any supported types according to the casting rule ''safe''