DOC: replace np.ma functions' return types with `MaskedArray` (#18964)

Relevant issue: #16751

Co-authored-by: paxcodes <13646646+paxcodes@users.noreply.github.com>
Co-authored-by: IrinaMaria <28827042+IrinaMaria@users.noreply.github.com>

1 files changed, 77 insertions, 11 deletions

diff --git a/numpy/ma/core.py b/numpy/ma/core.py
index 4c204cac2..82e5e7155 100644
--- a/numpy/ma/core.py
+++ b/numpy/ma/core.py
@@ -8094,22 +8094,51 @@ class _convert2ma:
     """
     __doc__ = None
 
-    def __init__(self, funcname, params=None):
+    def __init__(self, funcname, np_ret, np_ma_ret, params=None):
         self._func = getattr(np, funcname)
-        self.__doc__ = self.getdoc()
+        self.__doc__ = self.getdoc(np_ret, np_ma_ret)
         self._extras = params or {}
 
-    def getdoc(self):
+    def getdoc(self, np_ret, np_ma_ret):
         "Return the doc of the function (from the doc of the method)."
         doc = getattr(self._func, '__doc__', None)
         sig = get_object_signature(self._func)
         if doc:
+            doc = self._replace_return_type(doc, np_ret, np_ma_ret)
             # Add the signature of the function at the beginning of the doc
             if sig:
                 sig = "%s%s\n" % (self._func.__name__, sig)
             doc = sig + doc
         return doc
 
+    def _replace_return_type(self, doc, np_ret, np_ma_ret):
+        """
+        Replace documentation of ``np`` function's return type.
+
+        Replaces it with the proper type for the ``np.ma`` function.
+
+        Parameters
+        ----------
+        doc : str
+            The documentation of the ``np`` method.
+        np_ret : str
+            The return type string of the ``np`` method that we want to
+            replace. (e.g. "out : ndarray")
+        np_ma_ret : str
+            The return type string of the ``np.ma`` method.
+            (e.g. "out : MaskedArray")
+        """
+        if np_ret not in doc:
+            raise RuntimeError(
+                f"Failed to replace `{np_ret}` with `{np_ma_ret}`. "
+                f"The documentation string for return type, {np_ret}, is not "
+                f"found in the docstring for `np.{self._func.__name__}`. "
+                f"Fix the docstring for `np.{self._func.__name__}` or "
+                "update the expected string for return type."
+            )
+
+        return doc.replace(np_ret, np_ma_ret)
+
     def __call__(self, *args, **params):
         # Find the common parameters to the call and the definition
         _extras = self._extras
@@ -8125,20 +8154,57 @@ class _convert2ma:
             result._hardmask = bool(_extras.get("hard_mask", False))
         return result
 
-arange = _convert2ma('arange', params=dict(fill_value=None, hardmask=False))
+
+arange = _convert2ma(
+    'arange', 
+    params=dict(fill_value=None, hardmask=False),
+    np_ret='arange : ndarray',
+    np_ma_ret='arange : MaskedArray',
+)
 clip = np.clip
 diff = np.diff
-empty = _convert2ma('empty', params=dict(fill_value=None, hardmask=False))
-empty_like = _convert2ma('empty_like')
-frombuffer = _convert2ma('frombuffer')
-fromfunction = _convert2ma('fromfunction')
+empty = _convert2ma(
+    'empty', 
+    params=dict(fill_value=None, hardmask=False),
+    np_ret='out : ndarray',
+    np_ma_ret='out : MaskedArray',
+)
+empty_like = _convert2ma(
+    'empty_like',
+    np_ret='out : ndarray',
+    np_ma_ret='out : MaskedArray',
+)
+frombuffer = _convert2ma(
+    'frombuffer',
+    np_ret='out : ndarray',
+    np_ma_ret='out: MaskedArray',
+)
+fromfunction = _convert2ma(
+   'fromfunction',
+   np_ret='fromfunction : any',
+   np_ma_ret='fromfunction: MaskedArray',
+)
 identity = _convert2ma(
-    'identity', params=dict(fill_value=None, hardmask=False))
+    'identity', 
+    params=dict(fill_value=None, hardmask=False),
+    np_ret='out : ndarray',
+    np_ma_ret='out : MaskedArray',
+)
 indices = np.indices
-ones = _convert2ma('ones', params=dict(fill_value=None, hardmask=False))
+ones = _convert2ma(
+    'ones',
+    params=dict(fill_value=None, hardmask=False),
+    np_ret='out : ndarray',
+    np_ma_ret='out : MaskedArray',
+)
 ones_like = np.ones_like
 squeeze = np.squeeze
-zeros = _convert2ma('zeros', params=dict(fill_value=None, hardmask=False))
+zeros = _convert2ma(
+    'zeros',
+    params=dict(fill_value=None, hardmask=False),
+    np_ret='out : ndarray',
+    np_ma_ret='out : MaskedArray',
+)
 zeros_like = np.zeros_like