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]

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".