NumPyOverloader#

class overload_numpy.NumPyOverloader

Bases: Mapping[str, Dispatcher[Any]]

This mapping works in conjunction with a mixin (NPArrayFuncOverloadMixin, NPArrayUFuncOverloadMixin, or NPArrayOverloadMixin) to register and implement overrides with __array_function__ and __array_ufunc__.

Examples

First, some imports:

>>> from dataclasses import dataclass, fields
>>> from typing import ClassVar
>>> import numpy as np
>>> from overload_numpy import NumPyOverloader, NPArrayFuncOverloadMixin

Now we can define a NumPyOverloader instance:

>>> W_FUNCS = NumPyOverloader()

The overloads apply to an array wrapping class. Let’s define one:

>>> @dataclass
... class Wrap1D(NPArrayFuncOverloadMixin):
...     '''A simple array wrapper.'''
...     x: np.ndarray
...     NP_OVERLOADS: ClassVar[NumPyOverloader] = W_FUNCS

Now numpy functions can be overloaded and registered for Wrap1D.

>>> @W_FUNCS.implements(np.concatenate, Wrap1D)
... def concatenate(vecs):
...     VT = type(vecs[0])
...     return VT(*(np.concatenate(tuple(getattr(v, f.name) for v in vecs))
...                 for f in fields(VT)))

Time to check this works:

>>> vec1d = Wrap1D(np.arange(3))
>>> np.concatenate((vec1d, vec1d))
Wrap1D(x=array([0, 1, 2, 0, 1, 2]))

Methods

`assists`()	Register an `__array_function__` assistance function.
`get`(k[,d])
`implements`()	Register an `__array_function__` implementation object.
`items`()	Return registry items (str, dispatchers).
`keys`()	Return registry keys (`str`).
`values`()	Return registry values (dispatchers).

Initialize for dataclass-decorated subclasses.

Methods

`assists`()	Register an `__array_function__` assistance function.
`get`(k[,d])
`implements`()	Register an `__array_function__` implementation object.
`items`()	Return registry items (str, dispatchers).
`keys`()	Return registry keys (`str`).
`values`()	Return registry values (dispatchers).

Methods Summary

`assists`()	Register an `__array_function__` assistance function.
`implements`()	Register an `__array_function__` implementation object.
`items`()	Return registry items (str, dispatchers).
`keys`()	Return registry keys (`str`).
`values`()	Return registry values (dispatchers).

Methods Documentation

assists(numpy_funcs: UFuncLike, /, dispatch_on: type, *, types: type | TypeConstraint | Collection[type | TypeConstraint] | None, methods: Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat'] | Set[Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat']]) → OverloadUFuncDecorator[AssistsUFunc]#

assists(numpy_funcs: FunctionType, /, dispatch_on: type, *, types: type | TypeConstraint | Collection[type | TypeConstraint] | None, methods: Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat'] | Set[Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat']]) → OverloadFuncDecorator[AssistsFunc]

assists(numpy_funcs: set[Callable[..., Any] | UFuncLike], /, dispatch_on: type, *, types: type | TypeConstraint | Collection[type | TypeConstraint] | None = None, methods: Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat'] | Set[Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat']] = '__call__') → AssistsManyDecorator

This is a decorator factory, returning decorator, which registers the decorated function as an overload method for numpy function numpy_func for a class of type dispatch_on.

Parameters:

assistspython:callable()[…, Any], positional-only: The numpy function(s) that is(/are) being overloaded.
dispatch_ontype: The class type for which the overload implementation is being registered.
typestype or TypeConstraint or Collection thereof or python:None, keyword-only: The types of the arguments of assists. See __array_function__ for details. Only used if a function (not ufunc) is being overridden. If None then dispatch_on must have class-level attribute NP_FUNC_TYPES specifying the types.
methods{‘__call__’, ‘at’, ‘accumulate’, ‘outer’, ‘reduce’, ‘reduceat’} or set thereof, keyword-only: The ufunc methods for which this override applies. Default is just “__call__”. Only used if a ufunc (not function) is being overridden.

Returns:

AssistsManyDecorator: Decorator to register the wrapped function(s). This decorator should never be called by the user.

Examples

There’s a fair bit of setup required:

>>> from dataclasses import dataclass, fields
>>> from typing import ClassVar
>>> import numpy as np
>>> from overload_numpy import NumPyOverloader, NPArrayFuncOverloadMixin

>>> W_FUNCS = NumPyOverloader()

>>> @dataclass
... class Wrap1D(NPArrayFuncOverloadMixin):
...     '''A simple array wrapper.'''
...     x: np.ndarray
...     NP_OVERLOADS: ClassVar[NumPyOverloader] = W_FUNCS

>>> w1d = Wrap1D(np.arange(3))

Now we can register assists functions.

>>> stack_funcs = {np.vstack, np.hstack, np.dstack, np.column_stack, np.row_stack}
>>> @W_FUNCS.assists(stack_funcs, types=Wrap1D, dispatch_on=Wrap1D)
... def stack_assists(dispatch_on, func, vecs, *args, **kwargs):
...     cls = type(vecs[0])
...     return cls(*(func(tuple(getattr(v, f.name) for v in vecs), *args, **kwargs)
...                     for f in fields(cls)))

Checking this works:

>>> np.vstack((w1d, w1d))
Wrap1D(x=array([[0, 1, 2],
                [0, 1, 2]]))

>>> np.hstack((w1d, w1d))
Wrap1D(x=array([0, 1, 2, 0, 1, 2]))

assists also works for ufunc:

>>> add_funcs = {np.add, np.subtract}
>>> @W_FUNCS.assists(add_funcs, types=Wrap1D, dispatch_on=Wrap1D)
... def add_assists(cls, func, w1, w2, *args, **kwargs):
...     return cls(*(func(getattr(w1, f.name), getattr(w2, f.name), *args, **kwargs)
...                  for f in fields(cls)))

Checking this works:

>>> np.subtract(w1d, w1d)
Wrap1D(x=array([0, 0, 0]))

We can also to implement the ufunc methods, like``accumulate``, for all the add_funcs overloads:

>>> @add_assists.register("accumulate")
... def add_accumulate_assists(cls, func, w1, *args, **kwargs):
...     return cls(*(func(getattr(w1, f.name), *args, **kwargs)
...                  for f in fields(cls)))

>>> np.subtract.accumulate(w1d)
Wrap1D(x=array([ 0, -1, -3]))

implements(implements: UFuncLike, /, dispatch_on: type, *, types: type | TypeConstraint | Collection[type | TypeConstraint] | None = None, methods: Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat'] | Set[Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat']] = '__call__') → OverloadUFuncDecorator[ImplementsUFunc]#

implements(implements: FunctionType, /, dispatch_on: type, *, types: type | TypeConstraint | Collection[type | TypeConstraint] | None = None, methods: Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat'] | Set[Literal['__call__', 'at', 'accumulate', 'outer', 'reduce', 'reduceat']] = '__call__') → OverloadFuncDecorator[ImplementsFunc]

This is a decorator factory, returning decorator, which registers the decorated function as an overload method for numpy function implements for a class of type dispatch_on.

Parameters:

implementspython:callable()[…, Any], positional-only: The numpy function that is being overloaded.
dispatch_ontype: The class type for which the overload implementation is being registered.
typestype or TypeConstraint or Collection thereof or python:None, keyword-only: The types of the arguments of implements. See __array_function__. If None then dispatch_on must have class-level attribute NP_FUNC_TYPES specifying the types.
methods{‘__call__’, ‘accumulate’, ‘outer’, ‘reduce’} or python:None, keyword-only: numpy.ufunc methods.

Returns:

ImplementsFuncDecorator: Decorator to register the wrapped function.

Examples

There’s a fair bit of setup required:

>>> from dataclasses import dataclass, fields
>>> from typing import ClassVar
>>> import numpy as np
>>> from overload_numpy import NumPyOverloader, NPArrayFuncOverloadMixin

>>> W_FUNCS = NumPyOverloader()

>>> @dataclass
... class Wrap1D(NPArrayFuncOverloadMixin):
...     '''A simple array wrapper.'''
...     x: np.ndarray
...     NP_OVERLOADS: ClassVar[NumPyOverloader] = W_FUNCS

>>> w1d = Wrap1D(np.arange(3))

Now we can register an implements functions.

>>> @W_FUNCS.implements(np.concatenate, Wrap1D)  # overriding
... def concatenate(vecs):
...     VT = type(vecs[0])
...     return VT(*(np.concatenate(tuple(getattr(v, f.name) for v in vecs))
...                 for f in fields(VT)))

Checking it works:

>>> vec1d = Wrap1D(np.arange(3))

>>> newvec = np.concatenate((vec1d, vec1d))
>>> newvec
Vector2D(x=array([0, 1, 2, 0, 1, 2]), y=array([3, 4, 5, 3, 4, 5]))

implements also works for ufunc:

>>> @W_FUNCS.implements(np.add, Wrap1D)
... def add(w1, w2):
...     return Wrap1D(np.add(w1.x, w2.x))

>>> np.add(w1d, w1d)
Wrap1D(x=array([0, 2, 4]))

items() → ItemsView[str, All_Dispatchers]#: Return registry items (str, dispatchers).

keys() → KeysView[str]#: Return registry keys (str).

values() → ValuesView[All_Dispatchers]#: Return registry values (dispatchers).