Type Constraints (overload_numpy.constraints)

Classes for defining type constraints in __array_function__.

__array_function__ has an argument types, which is a Collection of unique argument types from the original NumPy function call that implement __array_function__. The purpose of types is to allow implementations of __array_function__ to check if all arguments of a type that the overload knows how to handle. Normally this is implemented inside of __array_function__, but overload_numpy gives overloading functions more flexibility to set constrains on a per-overloaded function basis.

Examples

First, some imports:

>>> from dataclasses import dataclass
>>> 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, and a subclass that contains more information:

>>> @dataclass
... class Wrap1D(NPArrayFuncOverloadMixin):
...     NP_OVERLOADS: ClassVar[NumPyOverloader] = W_FUNCS
...     NP_FUNC_TYPES: ClassVar[None] = None
...     x: np.ndarray

Note the NP_FUNC_TYPES. Normally this would inherit from NPArrayFuncOverloadMixin and be an empty frozenset, signalling that types are covariant with Wrap1D (includes subclasses). Changing NP_FUNC_TYPES to None means each overload must explicitly define the types argument. See NPArrayFuncOverloadMixin for further details.

Now numpy functions can be overloaded and registered for Wrap1D. Here is where we introduce the type constraints:

>>> from overload_numpy.constraints import Invariant, Covariant
>>> invariant = Invariant(Wrap1D)  # only works for Wrap1D

>>> @W_FUNCS.implements(np.concatenate, Wrap1D, types=invariant)
... def concatenate(w1ds):
...     return Wrap1D(np.concatenate(tuple(w.x for w in w1ds)))

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

This implementation of numpy.concatenate is invariant on the type of Wrap1D, so should fail if we use a subclass:

>>> @dataclass
... class Wrap2D(Wrap1D):
...     '''A simple 2-array wrapper.'''
...     y: np.ndarray

>>> w2d = Wrap2D(np.arange(3), np.arange(3, 6))
>>> try: np.concatenate((w2d, w2d))
>>> except Exception as e: print(e)
there is no implementation

Normally overlaoded functions are made to be Covariant. As this is the default, just passing the type is a convenience short-hand.

>>> @W_FUNCS.implements(np.concatenate, Wrap1D, types=Wrap1D)
... def concatenate(w1ds):
...     return Wrap1D(*(np.concatenate(tuple(getattr(w, f.name) for w in w1ds))
                        for f in fields(wlds[0])))

>>> np.concatenate((w2d, w2d))
Wrap1D(x=array([0, 1, 2, 0, 1, 2]), y=array([3, 4, 5, 3, 4, 5]))

This module offers other types of constraints, so be sure to check them out. Also, if you need something more specific, it’s easy to make your own constraint. There are currently two things you need to do:

1. subclass overload_numpy.constraints.TypeConstraint

2. define a method validate_type

As an example, let’s define a constraint where the argument must be one of 2 types:

>>> from overload_numpy.constraints import TypeConstraint
>>> @dataclass(frozen=True)
... class ThisOrThat(TypeConstraint):
...     this: type
...     that: type
...     def validate_type(self, arg_type: type, /) -> bool:
...         return arg_type is self.this or arg_type is self.that

Note

TypeConstraint will eventually be converted to a runtime-checkable typing.Protocol. When that happens step 1 (subclassing TypeConstraint) will become optional.

API

overload_numpy.constraints Module

Classes

TypeConstraint()	ABC for constraining an argument type.
Invariant(bound)	Type constraint for invariance -- the exact type.
Covariant(bound)	A covariant constraint -- permitting subclasses.
Contravariant(bound)	A contravariant constraint -- permitting superclasses.
Between(lower_bound, upper_bound)	Type constrained between two types.

Class Inheritance Diagram