Do not trigger type_of on return type hint

I am working with nested NumPy arrays (arrays of arrays) and I want to dispatch based on its nesting level.

from plum import dispatch

def something(a: NestedArray[0]):
    pass

def something(a: NestedArray[1]):
    pass

I’m using the following hook for that:

from plum import parametric, type_of

@parametric(runtime_type_of=True)
class NestedArray(np.ndarray):
    """A type for recursive numpy arrays (array of arrays) where the type parameter specifies the nesting level."""
    pass


@type_of.dispatch
def type_of(x: np.ndarray):
    level = 0
    while isinstance(x.flat[0], np.ndarray):
        level += 1
        x = x.flat[0]

    return NestedArray[level]

It works like a charm. But I run into a really annoying side-effect when I have a dispatched function with numpy.ndarray as a return type hint:

@dispatch
def to_numpy(a) -> numpy.ndarray:
    ...
    return numpy.zeros(4) # returns a ndarray, zeros for example

This code crashes because it runs type_of on the returning object, which returns a Type. Then compares it against the type hint of to_numpy, which fails.

...
  File "plum/function.py", line 537, in plum.function.Function.__call__
  File "plum/function.py", line 173, in plum.function._convert
  File "plum/function.py", line 537, in plum.function.Function.__call__
  File "/home/mofeing/Develop/rosnet/.venv/lib/python3.8/site-packages/plum/promotion.py", line 32, in convert
    return _convert.invoke(type_of(obj), type_to)(obj, type_to)
  File "plum/function.py", line 552, in plum.function.Function.invoke.wrapped_method
  File "/home/mofeing/Develop/rosnet/.venv/lib/python3.8/site-packages/plum/promotion.py", line 43, in _convert
    if type_from <= type_to:
  File "/home/mofeing/Develop/rosnet/.venv/lib/python3.8/site-packages/plum/util.py", line 43, in __ge__
    return other.__le__(self)
TypeError: descriptor '__le__' requires a 'numpy.ndarray' object but received a 'Type'

If I remove the return type hint, it works.

I guess this is a bug because there is no way (yet) to dispatch based on return type. Also because the function should be returning the object, not the hooked type.

[wesselb]

@mofeing I’m glad to hear that! I’m leaving this issue open until the bug has been fixed.

[mofeing]

@wesselb it works now with the call to ptype! thanks