Here's the code: https://github.com/rmcgibbo/numpy-mypy.
To implement it, I have the beginnings of the stubs that you'd expect, borrowed from
https://github.com/machinalis/mypy-data and then revised. Then, on top of that, I wrote some special type-level functions that are implemented inside of a mypy plugin. So, for example,
the stub's signature for np.sum is
|
def sum(a: ndarray[_S, _D], axis: AxesType=None, dtype: DtypeType=None, out: ndarray=None, keepdims: bool=False) -> ndarray[_InferDtypeWithDefault[_S], _InferNdimsReduction[_D]]: ... |
|
When the stub is applied, the resut's dtype is determined application of the _InferDtypeWithDefault type function, which defaults, as expected, to the dtype of the input array but checks of that was overridden dtype=None kwarg as well. And the _InferNdimsReduction type function has to check the axis and keepdims arguments as well.
It's by no means ready for real users, but I hope this is a useful place to build from. Any feedback or contributions would be appreciated.
-Robert