Proof of concept: annotation-based argument extractor/replacer

[peterbell10]

This explores the posibility of automatically generating the argument extractor and replacer based on inline type annotations (for simple cases). So, for example scipy.fft.fft might look like:

@_dispatch
def fft(x: ua.typing.Dispatchable[ArrayLike, np.ndarray], n=None, axis=-1,
        norm=None, overwrite_x=False, workers=None, *, plan=None):
    pass

The signature is inspected ahead-of-time so as to minimize call overhead but a Python implementation is still too slow. ~400 ns for
even a simple f(x) signature. However, if it were lowered to C++ then the concept may still be viable.

This also relies on typing.Annotated to specify dispatch_type and coercible but typing.Annotated itself requires Python 3.9. It may be possible to emulate on earlier python versions, otherwise I think the dispatch metadata would need to go elsewhere. Perhaps something like:

@_dispatch([ua.DispatchableArg(name='x', dispatch_type=np.ndarray, coercible=True)])
def fft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, plan=None):
    pass

cc @thomasjpfan

[hameerabbasi]

I'd be very supportive of this! It'd cover almost all use-cases besides np.block or something crazy like that. I think we could eventually even do np.concatenate and np.stack.

[BvB93]

@_dispatch
def fft(x: ua.typing.Dispatchable[ArrayLike, np.ndarray], n=None, axis=-1,
       norm=None, overwrite_x=False, workers=None, *, plan=None):
  pass

Should ArrayLike be interpretted as the annotation for x here by type checkers? If so, you may wanna trick them into treating ua.typing.Dispatchable as an alias for typing.Annotated, otherwise there is a big chance type checkers won't be able to understand the necasary logic behind __class_getitem__ (the heavy special casing behind typing.Annotated is not exactly helpful here either).

Also, something that might be important to remember is from __future__ import annotations, which converts all annotations into strings within a given module.

__all__ = ["Dispatchable"]

if not TYPE_CHECKING:
    class Dispatchable: ...  # Insert class details
else:
    from typing import Annotated as Dispatchable

def test_func(a: Dispatchable[str | list[str], str], b: int) -> None: ...

# Revealed type is "def (a: Union[builtins.str, builtins.list[builtins.str]], b: builtins.int)"
reveal_type(test_func)

[peterbell10]

otherwise there is a big chance type checkers won't be able to understand the necasary logic behind __class_getitem__

Hrm you're right. It seems mypy doesn't actually evaluate the types at all, so __class_getitem__ is never run (see python/mypy#11501). Instead it special cases classes that inherit from Generic[T] which doesn't work here because coercible isn't a type.

In that case I'm thinking DispatchableArg is the way to go. It's compatible with all Python versions and isn't really any longer than writing Annotated[T, Dispatchable(...)].

[thomasjpfan]

In general, I like the idea of using a decorator because it simplifies the workflow for API authors.

Minor nit: If we are going with DispatchableArg, then I prefer this API:

@_dispatch(x=ua.DispatchableArg(type=np.ndarray, coercible=True))
def fft(x, n=None, axis=-1, norm=None, overwrite_x=False, workers=None, *, plan=None):
    pass

[peterbell10]

@thomasjpfan I'm not sure that will work because the general decorator also needs to specify domain. So what if your function has an argument called domain. We could support dict(x=ua.DispatchableArg(...)) though I suppose. Not sure if that's actually cleaner though.

[thomasjpfan]

Brainstorming: A dispatcher for a given domain:

dispatch = Dispatcher(domain="scipy.fft")

@dispatch.register(
	x=DispatchableArg(...)
)
def fft(x):
    ...

I think x= is slightly cleaner because x= looks more like a parameter. With there is some name="x" minor cognitive overhead connecting the string "x" to the function parameter x.

Edit: I see other issues with using kwargs. Likely name="x" is best.