plum
plum copied to clipboard
beartype
Plum 2 and np.typing.NDArray
Hi @wesselb
I am testing out https://github.com/wesselb/plum/pull/73 due to the improved numpy support. In the very first example there, you show how to dispatch based on shape and type. In my use cases, I am only interested in dispatching based on type. Since that is actually supported by np.typing.NDArray, I thought I could just use that, but that doesn't seem to be the case.
I report a minimal example below, in case you want to have a look before the plum 2 release.
from plum import dispatch, parametric
from typing import Any, Optional, Tuple, Union
import numpy as np
import numpy.typing
class NDArrayMeta(type):
def __instancecheck__(self, x):
if self.concrete:
shape, dtype = self.type_parameter
else:
shape, dtype = None, None
return (
isinstance(x, np.ndarray)
and (shape is None or x.shape == shape)
and (dtype is None or x.dtype == dtype)
)
@parametric
class NDArray(np.ndarray, metaclass=NDArrayMeta):
@classmethod
@dispatch
def __init_type_parameter__(
cls,
shape: Optional[Tuple[int, ...]],
dtype: Optional[Any],
):
"""Validate the type parameter."""
return shape, dtype
@classmethod
@dispatch
def __le_type_parameter__(
cls,
left: Tuple[Optional[Tuple[int, ...]], Optional[Any]],
right: Tuple[Optional[Tuple[int, ...]], Optional[Any]],
):
"""Define an order on type parameters. That is, check whether
`left <= right` or not."""
shape_left, dtype_left = left
shape_right, dtype_right = right
return (
(shape_right is None or shape_left == shape_right)
and (dtype_right is None or dtype_left == dtype_right)
)
@dispatch
def f(x: NDArray[None, np.int32]):
print("An int array!")
@dispatch
def f(x: NDArray[None, np.float64]):
print("A float array!")
print("BEGIN f")
f(np.ones((3, 3), np.int32))
f(np.ones((2, 2), np.float64))
print("END f")
@dispatch
def g(x: np.typing.NDArray[np.int32]):
print("An int array!")
@dispatch
def g(x: np.typing.NDArray[np.float64]):
print("A float array!")
print("BEGIN g")
g(np.ones((3, 3), np.int32))
g(np.ones((2, 2), np.float64))
print("END g")
and the corresponding output
BEGIN f
An int array!
A float array!
END f
BEGIN g
/usr/local/lib/python3.11/dist-packages/plum/signature.py:203: UserWarning: Could not resolve the type hint of `numpy.ndarray[typing.Any, numpy.dtype[numpy.int32]]`. I have ended the resolution here to not make your code break, but some types might not be working correctly. Please open an issue at https://github.com/wesselb/plum.
annotation = resolve_type_hint(p.annotation)
/usr/local/lib/python3.11/dist-packages/plum/type.py:261: UserWarning: Could not resolve the type hint of `numpy.ndarray[typing.Any, numpy.dtype[numpy.int32]]`. I have ended the resolution here to not make your code break, but some types might not be working correctly. Please open an issue at https://github.com/wesselb/plum.
return _is_faithful(resolve_type_hint(x))
/usr/local/lib/python3.11/dist-packages/plum/type.py:261: UserWarning: Could not determine whether `numpy.ndarray[typing.Any, numpy.dtype[numpy.int32]]` is faithful or not. I have concluded that the type is not faithful, so your code might run with subpar performance. Please open an issue at https://github.com/wesselb/plum.
return _is_faithful(resolve_type_hint(x))
/usr/local/lib/python3.11/dist-packages/plum/signature.py:203: UserWarning: Could not resolve the type hint of `numpy.ndarray[typing.Any, numpy.dtype[numpy.float64]]`. I have ended the resolution here to not make your code break, but some types might not be working correctly. Please open an issue at https://github.com/wesselb/plum.
annotation = resolve_type_hint(p.annotation)
/usr/local/lib/python3.11/dist-packages/plum/type.py:261: UserWarning: Could not resolve the type hint of `numpy.ndarray[typing.Any, numpy.dtype[numpy.float64]]`. I have ended the resolution here to not make your code break, but some types might not be working correctly. Please open an issue at https://github.com/wesselb/plum.
return _is_faithful(resolve_type_hint(x))
/usr/local/lib/python3.11/dist-packages/plum/type.py:261: UserWarning: Could not determine whether `numpy.ndarray[typing.Any, numpy.dtype[numpy.float64]]` is faithful or not. I have concluded that the type is not faithful, so your code might run with subpar performance. Please open an issue at https://github.com/wesselb/plum.
return _is_faithful(resolve_type_hint(x))
Traceback (most recent call last):
File "/tmp/p.py", line 72, in <module>
g(np.ones((3, 3), np.int32))
File "/usr/local/lib/python3.11/dist-packages/plum/function.py", line 342, in __call__
self._resolve_pending_registrations()
File "/usr/local/lib/python3.11/dist-packages/plum/function.py", line 237, in _resolve_pending_registrations
self._resolver.register(subsignature)
File "/usr/local/lib/python3.11/dist-packages/plum/resolver.py", line 58, in register
existing = [s == signature for s in self.signatures]
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/plum/resolver.py", line 58, in <listcomp>
existing = [s == signature for s in self.signatures]
^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/plum/util.py", line 132, in __eq__
return self <= other <= self
^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/plum/signature.py", line 132, in __le__
[TypeHint(x) <= TypeHint(y) for x, y in zip(self_types, other_types)]
File "/usr/local/lib/python3.11/dist-packages/plum/signature.py", line 132, in <listcomp>
[TypeHint(x) <= TypeHint(y) for x, y in zip(self_types, other_types)]
^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/beartype/door/_doormeta.py", line 148, in __call__
_HINT_KEY_TO_WRAPPER.cache_or_get_cached_func_return_passed_arg(
File "/usr/local/lib/python3.11/dist-packages/beartype/_util/cache/map/utilmapbig.py", line 231, in cache_or_get_cached_func_return_passed_arg
value = value_factory(arg)
^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/beartype/door/_doormeta.py", line 220, in _make_wrapper
wrapper_subclass = get_typehint_subclass(hint)
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File "/usr/local/lib/python3.11/dist-packages/beartype/door/_doordata.py", line 108, in get_typehint_subclass
raise BeartypeDoorNonpepException(
beartype.roar.BeartypeDoorNonpepException: Type hint numpy.ndarray[typing.Any, numpy.dtype[numpy.int32]] invalid (i.e., either PEP-noncompliant or PEP-compliant but currently unsupported by "beartype.door.TypeHint").
Thanks!
Hey @francesco-ballarin,
Thanks for opening an issue about this. :) You're right that numpy.typing currently won't work. The problem is that numpy.typing types are unfortunately not functional themselves:
>>> isinstance(1, npt.NDArray[int])
TypeError: isinstance() argument 2 cannot be a parameterized generic
They are type hints like objects from typing, which won't of their own, but need additional support. @beartype seems to come very close, but unfortunately doesn't quite get it right:
>>> from beartype.door import is_bearable, TypeHint
>>> is_bearable(np.ones(1, int), npt.NDArray[int]) # Nice!
True
>>> TypeHint(npt.NDArray[int]) # Nice!
TypeHint(numpy.ndarray[typing.Any, numpy.dtype[int]])
>>> TypeHint(npt.NDArray[int]) == TypeHint(npt.NDArray[float]) # :(
TypeHint(npt.NDArray[int]) == TypeHint(npt.NDArray[float])
The reason that you're getting a beartype.roar.BeartypeDoorNonpepException whereas I'm not might be due to a different Python version. (I ran the above using Python 3.9.) EDIT: This seems to be the case; see the issue linked below.
I've opened an issue on @beartype to see what @leycec thinks about this.
