ENH: add mypy plugin for more precise ufunc signatures

[person142]

Keeping this as a draft as it needs discussion (if we decide its ok I'll need to add more data on ufunc arities to the plugin).

This adds a mypy plugin that alters ufunc signatures to have the correct arity at type checking time. The tl;dr version is:

np.sin(1, 1)  # Ok on master, but fails with this plugin

The broader issue is:

• ufuncs are all instances of the same class
• but their arities differ
• there's no way to capture that in the current type system

To work around that, the plugin lets us halt each time a numpy.ufunc.__call__ node is reached in the AST and modify the signature. This just does arities, but you can imagine that when ndarray is generic over dtypes then we can potentially do even fancier things like adding overloads for all the appropriate dispatches for the ufunc.

Now, I imagine one concern people might have is "doesn't this lock us into mypy", and the answer is "not really". The plugin is purely opt in-you have to include it in your mypy.ini for it to be activated. So we can easily have the stubs, which describe things as best they can, and a set of plugins that give more precise types for people who happen to be using mypy.

Thoughts?

[person142]

The punchline: we statically know that sin takes one argument.

[person142]

We have our own tests opt into using the plugin.

[person142]

We would need a big dictionary of the ufuncs and their arities.

[seberg]

UFuncs should provide all the information you need for this, no? Such as np.sin.nin, np.sin.nout, and np.sin.signature (which is None for normal ufuncs).

[person142]

I've been assuming it's bad practice to import NumPy instead of working just from the AST, but heck I don't know it seems pretty reasonable here... maybe we should just do it? Worried about triggering https://xkcd.com/292/.

[seberg]

Ohh, I didn't think about that and I guess it is a mypy/typing code smell question and I have no nose for typing :). I suppose we thus need something like a decorator for ufuncs that duplicates some of the information? (Its not like we can get the decorator information to the C-level cleanly probably).
Also, it is not like NumPy is the only one with ufuncs. SciPy will need to use this hook as well.

[person142]

Switched to pulling it off the ufunc object in f8c5979. One (and only one!) of the CI builds is failing though...

[person142]

Interestingly, PEP 612:

https://www.python.org/dev/peps/pep-0612/

might also allow us to better type ufuncs by making the ufunc class generic over a parameter specification.

[person142]

And here's a branch

https://github.com/person142/numpy-stubs/tree/ufunc-plugin-with-literal

with a plugin that abuses Literals to allow you to specify ufunc arity by doing something like

from typing_extensions import Literal

sin: ufunc[Literal[1], Literal[1]]
           ^           ^
           nin         nout

So I guess we've got some options:

• Include the ufunc arity data in the plugin itself
  • Pretty tame
  • Can't handle non-NumPy ufuncs
  • Not very DRY
• Import NumPy in the plugin
  • Maybe a little sketchy
  • Can't handle non-NumPy ufuncs
  • For some reason fails on 3.7 in CI (but works in 3.6 and 3.8, so might be a bug elsewhere)
• Abuse literals
  • Pretty out there
  • Can handle non-NumPy ufuncs

[seberg]

My main concern is that we really should have a solution that will work also for other libraries, and while importing numpy is possibly fine, we should maybe not force the library to do the same?

I am still curious if we cannot do the plugin solution anyway, but provide a way to "expand" the ufuncs, if necessary by telling projects how to create a very simply plugin themselves? The literal solution seems fine, although having Literal in there seems a bit ugly?

[person142]

Closing this; it was a good experiment with plugins (I think I feel more confident that they will be a big part of supporting types in NumPy), but it's not really "production grade".