def test_is_jitted(self):
def foo(x):
pass
self.assertFalse(is_jitted(foo))
self.assertTrue(is_jitted(njit(foo)))
self.assertFalse(is_jitted(vectorize(foo)))
self.assertFalse(is_jitted(vectorize(parallel=True)(foo)))
self.assertFalse(is_jitted(guvectorize("void(float64[:])", "(m)")(foo)))
def jit(function: TFunc,
signature=None,
parallel: bool = False,
**kwargs) -> TFunc:
"""apply nb.jit with predefined arguments
Args:
signature: Signature of the function to compile
parallel (bool): Allow parallel compilation of the function
**kwargs: Additional arguments to `nb.jit`
Returns:
Function that will be compiled using numba
"""
if is_jitted(function):
return function
# prepare the compilation arguments
kwargs.setdefault("nopython", True)
jit_kwargs = _numba_get_signature(parallel=parallel, **kwargs)
# log some details
logger = logging.getLogger(__name__)
name = function.__name__ # type: ignore
if kwargs["nopython"]: # standard case
logger.info("Compile `%s` with parallel=%s", name,
jit_kwargs["parallel"])
else: # this might imply numba falls back to object mode
logger.warning("Compile `%s` with nopython=False", name)
return nb.jit(signature, **jit_kwargs)(function) # type: ignore
def decorator(f):
@functools.wraps(f)
def wrapper(*args, **kwargs):
return f(*args, **kwargs)
__bound_lite_func__ = _LiteFuncDict()
attrs["lite"] = lite_func
for bound_name, attr in attrs.items():
# only allow functions or jitted functions
if not (inspect.isfunction(attr) or is_jitted(attr)):
raise ValueError(
f"Can not bind obj '{attr}' to function '{wrapper.__name__}'."
f" Only functions are allowed to be bound. Skipping."
)
# build origin name
if hasattr(attr, "__module__"):
modname = attr.__module__
else: # coverage: ignore
# assume attr is defined in the module the function being
# decorated in
modname = wrapper.__module__
origin = f"{modname}.{attr.__name__}"
__bound_lite_func__[bound_name] = origin
# bind
setattr(wrapper, bound_name, attr)
setattr(wrapper, "__bound_lite_func__", __bound_lite_func__)
return wrapper
def test_is_jitted(self):
"""Ensure `thermal_speed_lite` was jitted by `numba`."""
assert is_jitted(thermal_speed_lite)
def test_is_jitted(self):
"Ensure `plasmapy_frequency_lite` was jitted by `numba`."
assert is_jitted(plasma_frequency_lite)
def bind_lite_func(lite_func, attrs: Dict[str, Callable] = None):
"""
Decorator to bind a lightweight "lite" version of a formulary
function to the full formulary function, as well as any supporting
attributes.
Parameters
----------
lite_func: Callable
The lightweight function to be bound as the ``lite`` attribute
to the function being decorated.
attrs: Dict[str, Callable]
A dictionary where the key is a string defining the bound name
and the associated value is the functionality to be bound.
Examples
--------
.. code-block:: python
def foo_lite(x)
return x
def bar():
print("Supporting function.")
@bind_lite_func(foo_lite, attrs=[("bar", bar),])
def foo(x):
if not isinstance(x, float):
raise TypeError("Argument x can only be a float.")
return x
>>> foo(5) # doctest: +SKIP
5
>>> foo.lite(5) # doctest: +SKIP
5
>>> foo.bar() # doctest: +SKIP
Supporting function.
Notes
-----
In addition to binding the functionality defined by the inputs, a
``__bound_lite_func__`` dunder is bound. This dunder is a
dictionary where a key is a string representing the bound name of
the bound functionality and the associated value is a string
representing the fully qualified path of the original bound
functionality.
"""
if attrs is None:
attrs = {}
elif not isinstance(attrs, dict):
raise TypeError(
f"Argument 'attrs' is a type '{type(attrs)}', expected a dictionary."
)
elif "lite" in attrs:
raise ValueError(
"Argument 'attr' can NOT define key 'lite', this is reserved for"
" the 'lite_func' argument."
)
if inspect.isbuiltin(lite_func) or not (
is_jitted(lite_func) or inspect.isfunction(lite_func)
):
raise ValueError("The given lite-function is not a user-defined function.")
def decorator(f):
@functools.wraps(f)
def wrapper(*args, **kwargs):
return f(*args, **kwargs)
__bound_lite_func__ = _LiteFuncDict()
attrs["lite"] = lite_func
for bound_name, attr in attrs.items():
# only allow functions or jitted functions
if not (inspect.isfunction(attr) or is_jitted(attr)):
raise ValueError(
f"Can not bind obj '{attr}' to function '{wrapper.__name__}'."
f" Only functions are allowed to be bound. Skipping."
)
# build origin name
if hasattr(attr, "__module__"):
modname = attr.__module__
else: # coverage: ignore
# assume attr is defined in the module the function being
# decorated in
modname = wrapper.__module__
origin = f"{modname}.{attr.__name__}"
__bound_lite_func__[bound_name] = origin
# bind
setattr(wrapper, bound_name, attr)
setattr(wrapper, "__bound_lite_func__", __bound_lite_func__)
return wrapper
return decorator