ENH: add type annotations to numpy.ma

[johnthagen]

Describe the issue:

Running Mypy on some functions within numpy.ma lead to type checking errors due to missing annotations.

Reproduce the code example:

import numpy

a = numpy.ma.array([1, 2, 3, 4], mask=[0, 0, 1, 0])
b = numpy.ma.masked_array([1, 2, 3, 4], mask=[0, 0, 1, 0])

Error message:

$ mypy test.py 
test.py:3: error: Call to untyped function "array" in typed context  [no-untyped-call]
test.py:4: error: Call to untyped function "MaskedArray" in typed context  [no-untyped-call]

Python and NumPy Versions:

1.26.4
3.12.2 (main, Feb  6 2024, 20:19:44) [Clang 15.0.0 (clang-1500.1.0.2.5)]

Runtime Environment:

[{'numpy_version': '1.26.4',
  'python': '3.12.2 (main, Feb  6 2024, 20:19:44) [Clang 15.0.0 '
            '(clang-1500.1.0.2.5)]',
  'uname': uname_result(system='Darwin', node='M2DF94775W-ML', release='23.4.0', version='Darwin Kernel Version 23.4.0: Fri Mar 15 00:12:37 PDT 2024; root:xnu-10063.101.17~1/RELEASE_ARM64_T6031', machine='arm64')},
 {'simd_extensions': {'baseline': ['NEON', 'NEON_FP16', 'NEON_VFPV4', 'ASIMD'],
                      'found': ['ASIMDHP'],
                      'not_found': ['ASIMDFHM']}},
 {'architecture': 'armv8',
  'filepath': '/Users/hagenjt1/Library/Caches/pypoetry/virtualenvs/non-package-mode-eSjUyYda-py3.12/lib/python3.12/site-packages/numpy/.dylibs/libopenblas64_.0.dylib',
  'internal_api': 'openblas',
  'num_threads': 16,
  'prefix': 'libopenblas',
  'threading_layer': 'pthreads',
  'user_api': 'blas',
  'version': '0.3.23.dev'}]
None

Context for the issue:

We rely on type checking to find bugs in our software that uses Numpy before it is executed. This enables higher code quality and reliability. False positive type errors cost time to track down and can distract from real issues.

Type annotations also help IDEs such as PyCharm and VS Code perform better intellisense during development.

[rgommers]

Thanks @johnthagen. This is more of a feature request I think. All of numpy.ma has type stubs, but no return types. The following diff will make your example pass (note: the actually correct way of doing this is a bit different, look at the class ndarray annotations in numpy/__init__.pyi for how to annotate an array class):

diff --git a/numpy/ma/core.pyi b/numpy/ma/core.pyi
index d6cc0a782c..c3f27e1271 100644
--- a/numpy/ma/core.pyi
+++ b/numpy/ma/core.pyi
@@ -167,7 +167,7 @@ def __next__(self): ...
 
 class MaskedArray(ndarray[_ShapeType, _DType_co]):
     __array_priority__: Any
-    def __new__(cls, data=..., mask=..., dtype=..., copy=..., subok=..., ndmin=..., fill_value=..., keep_mask=..., hard_mask=..., shrink=..., order=...): ...
+    def __new__(cls, data=..., mask=..., dtype=..., copy=..., subok=..., ndmin=..., fill_value=..., keep_mask=..., hard_mask=..., shrink=..., order=...) -> 'MaskedArray': ...
     def __array_finalize__(self, obj): ...
     def __array_wrap__(self, obj, context=..., return_scalar=...): ...
     def view(self, dtype=..., type=..., fill_value=...): ...
@@ -358,7 +358,7 @@ def array(
     shrink=...,
     subok=...,
     ndmin=...,
-): ...
+) -> MaskedArray: ...
 def is_masked(x): ...
 
 class _extrema_operation(_MaskedUFunc):

But that whole core.pyi file needs to be expanded to have return types. That's not a small amount of work.

I'll change this to an enhancement request.

[johnthagen]

I'll change this to an enhancement request.

No problem! I was under the impression that all of Numpy fully supported static typing (and even includes its own Mypy plugin) so that it why I thought it was a bug report if some of the static typing was incorrect/missing.

Regardless, thank you for the clear feedback and consideration.

[rgommers]

I think all namespaces should have entries for all public objects, and the main namespaces and most important other namespaces (fft, linalg, etc.) are fully typed. But some other ones are not.

[luxedo]

I'd be glad to help! I Just need some guidance or a checklist. Is it literally all the functions and methods in ma/core.pyi? I could create a PR for MaskedArray or any smaller class first.

[rgommers]

Is it literally all the functions and methods in ma/core.pyi?

Pretty much, yes. Also the other numpy/ma/*.pyi files, but core.pyi contains the most important ones.

I'd suggest picking a small, representative subset and typing those. Once the review on those is done, that establishes the patterns to expand to the full set of functions/objects.

E.g.:

• functions: min, take, diag, ndim, set_fill_value, fix_invalid
• _MaskedUfunc
• MaskedArray methods/properties: all, copy, shape, mask

[rgommers]

And use the type annotations for the functions in the main namespace, as well as the np.ndarray/np.ufunc ones, for reference.

[luxedo]

Typing ma/core.pyi

The non-exhaustive list below serves as a guide for typing all of ma/core.pyi.

First batch

• functions
  • min
  • max
  • ptp
  • take
  • diag
  • ndim
  • shape
  • size
  • transpose
  • set_fill_value
  • fix_invalid
• _MaskedUfunc
• MaskedArray
  • all
  • copy
  • shape
  • mask

Second batch

• functions
  • power
  • argsort
  • sort
  • compressed
  • concatenate
  • diag
  • left_shift
  • right_shift
  • put
  • putmask
  • reshape
  • resize
  • diff
  • where
  • choose
  • round
• _convert2ma:
• _MaskedUnaryOperation
• _MaskedBinaryOperation
• MaskedArray
  • ...

Third batch

• functions
  • inner
  • outer
  • correlate
  • convolve
  • allequal
  • allclose
  • asarray
  • asanyarray
  • fromflex
  • append
  • dot
  • mask_rowcols
  • default_fill_value
  • minimum_fill_value
  • maximum_fill_value
  • common_fill_value
  • filled
  • getdata
• _DomainBinaryOperation
• _frommethod