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