ENH: implement NEP-35's `like=` argument (gh-16935)

This PR adds the implementation of NEP-35's like= argument, allowing dispatch of array creation functions with __array_function__ based on a reference array.

* ENH: Add like= kwarg via __array_function__ dispatcher to asarray

* ENH: Add new function for __array_function__ dispatching from C

This new function allows dispatching from C directly, while also
implementing the new `like=` argument, requiring only minimal
changes to existing array creation functions that need to add
support for that argument.

* ENH: Add like= support to numpy.array

The implementation uses array_implement_c_array_function, thus
introducing minimal complexity to the original _array_fromobject
code.

* BUG: Fix like= dispatcher for np.full

* ENH: Remove np.asarray like= dispatcher via Python

np.asarray can rely on np.array's C dispatcher instead.

* TST: Add some tests for like= argument

Tests comprise some of the functions that have been implemented already:
* np.array (C dispatcher)
* np.asarray (indirect C dispatcher via np.array)
* np.full (Python dispatcher)
* np.ones (Python dispatcher)

* ENH: Remove like= argument during array_implement_array_function

* ENH: Add like= kwarg to ones and full

* BUG: prevent duplicate removal of `like=` argument

* ENH: Make `like=` a keyword-only argument

* ENH: Use PyUnicode_InternFromString in arrayfunction_override

Replace PyUnicode_FromString by PyUnicode_InternFromString to cache
"like" string.

* ENH: Check for arrayfunction_override import errors

Check and handle errors on importing NumPy's Python functions

* BUG: Fix array_implement_c_array_function error handling

* ENH: Handle exceptions with C implementation of `like=`

* ENH: Add `like=` dispatch for all asarray functions

Using Python dispatcher for all of them. Using the C dispatcher
directly on the `np.array` call can result in incorrect behavior.
Incorrect behavior may happen if the downstream library's
implementation is different or if not all keyword arguments are
supported.

* ENH: Simplify handling of exceptions with `like=`

* TST: Add test for exception handling with `like=`

* ENH: Add support for `like=` to `np.empty` and `np.zeros`

* TST: Add `like=` tests for `np.empty` and `np.zeros`

* ENH: Add `like=` to remaining multiarraymodule.c functions

Functions are:
* np.arange
* np.frombuffer
* np.fromfile
* np.fromiter
* np.fromstring

* TST: Add tests for multiarraymodule.c functions with like=

Functions are:
* np.arange
* np.frombuffer
* np.fromfile
* np.fromiter
* np.fromstring

* ENH: Add `like=` support to more creation functions

Support for the following functions is added:
* np.eye
* np.fromfunction
* np.genfromtxt
* np.identity
* np.loadtxt
* np.tri

* TST: Add `like=` tests for multiple functions

Tests for the following functions are added:
* np.eye
* np.fromfunction
* np.genfromtxt
* np.identity
* np.loadtxt
* np.tri

* TST: Reduce code duplication in `like=` tests

* DOC: Document `like=` in functions that support it

Add documentations for the following functions:
* np.array
* np.arange
* np.asarray
* np.asanyarray
* np.ascontiguousarray
* np.asfortranarray
* np.require
* np.empty
* np.full
* np.ones
* np.zeros
* np.identity
* np.eye
* np.tri
* np.frombuffer
* np.fromfile
* np.fromiter
* np.fromstring
* np.loadtxt
* np.genfromtxt

* ENH: Add `like=` to numpy/__init__.pyi stubs

* BUG: Remove duplicate `like=` dispatching in as*array

Functions `np.asanyarray`, `np.contiguousarray` and `np.fortranarray`
were dispatching both via their definitions and `np.array` calls,
the latter should be avoided.

* BUG: Fix missing check in array_implement_array_function

* BUG: Add missing keyword-only markers in stubs

* BUG: Fix duplicate keyword-only marker in array stub

* BUG: Fix syntax error in numpy/__init__.pyi

* BUG: Fix more syntax errors in numpy/__init__.pyi

* ENH: Intern arrayfunction_override strings in multiarraymodule

* STY: Add missing brackets to arrayfunction_override.c

* MAINT: Remove arrayfunction_override dict check for kwarg

* TST: Assert that 'like' is not in TestArrayLike kwargs

* MAINT: Rename array_implement_c_array_function(_creation)

This is done to be more explicit as to its usage being intended for
array creation functions only.

* MAINT: Use NotImplemented to indicate fallback to default

* TST: Test that results with `like=np.array` are correct

* TST: Avoid duplicating MyArray code in TestArrayLike

* TST: Don't delete temp file, it may cause issues with Windows

* TST: Don't rely on eval in TestArrayLike

* TST: Use lambda with StringIO in TestArrayLike

* ENH: Avoid unnecessary Py_XDECREF in arrayfunction_override

* TST: Make TestArrayLike more readable

* ENH: Cleaner error handling in arrayfunction_override

* ENH: Simplify array_implement_c_array_function_creation

* STY: Add missing spaces to multiarraymodule.c

* STY: C99 declaration style in arrayfunction_override.c

* ENH: Simplify arrayfunction_override.c further

Remove cleanup label from array_implementation_c_array_function,
simplifying the code. Fix unitialized variable warning in
array_implementation_array_function_internal.

* DOC: Use string replacement for `like=` documentation

Avoid repeating the full text for the `like=` argument by storing it as
a variable and using `replace` on each docstring.

* DOC: Update `like=` docstring

* TST: Test like= with array not implementing __array_function__

* TST: Add missing asanyarray to TestArrayLike

* ENH: Use helper function for like= dispatching

Avoid dispatching like= from Python implementation functions to improve
their performance. This is achieved by only calling a dispatcher
function when like is passed by the users.

* ENH: Rename array_function_dispatch kwarg to public_api

* BUG: Add accidentally removed decorator for np.eye back

* DOC: Add set_array_function_like_doc function

The function keeps Python files cleaner and resolve errors when __doc__
is not defined due to PYTHONOPTIMIZE or -OO .

* DOC: Add mention to like= kwarg being experimental

* TST: Test like= with not implemented downstream function

* DOC: Fix like= docstring reference to NEP 35.

* ENH: Prevent silent errors if public_api is not callable

* ENH: Make set_array_function_like_doc a decorator

* ENH: Simplify `_*_with_like` functions

* BUG: Fix multiple `like=` dispatching in `require`

* MAINT: Remove now unused public_api from array_function_dispatch

Co-authored-by: Sebastian Berg <sebastian@sipsolutions.net>

1 files changed, 88 insertions, 6 deletions

diff --git a/numpy/core/_asarray.py b/numpy/core/_asarray.py
index 1b06c328f..a406308f3 100644
--- a/numpy/core/_asarray.py
+++ b/numpy/core/_asarray.py
@@ -3,7 +3,11 @@ Functions in the ``as*array`` family that promote array-likes into arrays.
 
 `require` fits this category despite its name not matching this pattern.
 """
-from .overrides import set_module
+from .overrides import (
+    array_function_dispatch,
+    set_array_function_like_doc,
+    set_module,
+)
 from .multiarray import array
 
 
@@ -11,8 +15,14 @@ __all__ = [
     "asarray", "asanyarray", "ascontiguousarray", "asfortranarray", "require",
 ]
 
+
+def _asarray_dispatcher(a, dtype=None, order=None, *, like=None):
+    return (like,)
+
+
+@set_array_function_like_doc
 @set_module('numpy')
-def asarray(a, dtype=None, order=None):
+def asarray(a, dtype=None, order=None, *, like=None):
     """Convert the input to an array.
 
     Parameters
@@ -30,6 +40,9 @@ def asarray(a, dtype=None, order=None):
         'A' (any) means 'F' if `a` is Fortran contiguous, 'C' otherwise
         'K' (keep) preserve input order
         Defaults to 'C'.
+    ${ARRAY_FUNCTION_LIKE}
+
+        .. versionadded:: 1.20.0
 
     Returns
     -------
@@ -83,11 +96,20 @@ def asarray(a, dtype=None, order=None):
     True
 
     """
+    if like is not None:
+        return _asarray_with_like(a, dtype=dtype, order=order, like=like)
+
     return array(a, dtype, copy=False, order=order)
 
 
+_asarray_with_like = array_function_dispatch(
+    _asarray_dispatcher
+)(asarray)
+
+
+@set_array_function_like_doc
 @set_module('numpy')
-def asanyarray(a, dtype=None, order=None):
+def asanyarray(a, dtype=None, order=None, *, like=None):
     """Convert the input to an ndarray, but pass ndarray subclasses through.
 
     Parameters
@@ -105,6 +127,9 @@ def asanyarray(a, dtype=None, order=None):
         'A' (any) means 'F' if `a` is Fortran contiguous, 'C' otherwise
         'K' (keep) preserve input order
         Defaults to 'C'.
+    ${ARRAY_FUNCTION_LIKE}
+
+        .. versionadded:: 1.20.0
 
     Returns
     -------
@@ -140,11 +165,24 @@ def asanyarray(a, dtype=None, order=None):
     True
 
     """
+    if like is not None:
+        return _asanyarray_with_like(a, dtype=dtype, order=order, like=like)
+
     return array(a, dtype, copy=False, order=order, subok=True)
 
 
+_asanyarray_with_like = array_function_dispatch(
+    _asarray_dispatcher
+)(asanyarray)
+
+
+def _asarray_contiguous_fortran_dispatcher(a, dtype=None, *, like=None):
+    return (like,)
+
+
+@set_array_function_like_doc
 @set_module('numpy')
-def ascontiguousarray(a, dtype=None):
+def ascontiguousarray(a, dtype=None, *, like=None):
     """
     Return a contiguous array (ndim >= 1) in memory (C order).
 
@@ -154,6 +192,9 @@ def ascontiguousarray(a, dtype=None):
         Input array.
     dtype : str or dtype object, optional
         Data-type of returned array.
+    ${ARRAY_FUNCTION_LIKE}
+
+        .. versionadded:: 1.20.0
 
     Returns
     -------
@@ -181,11 +222,20 @@ def ascontiguousarray(a, dtype=None):
     so it will not preserve 0-d arrays.  
 
     """
+    if like is not None:
+        return _ascontiguousarray_with_like(a, dtype=dtype, like=like)
+
     return array(a, dtype, copy=False, order='C', ndmin=1)
 
 
+_ascontiguousarray_with_like = array_function_dispatch(
+    _asarray_contiguous_fortran_dispatcher
+)(ascontiguousarray)
+
+
+@set_array_function_like_doc
 @set_module('numpy')
-def asfortranarray(a, dtype=None):
+def asfortranarray(a, dtype=None, *, like=None):
     """
     Return an array (ndim >= 1) laid out in Fortran order in memory.
 
@@ -195,6 +245,9 @@ def asfortranarray(a, dtype=None):
         Input array.
     dtype : str or dtype object, optional
         By default, the data-type is inferred from the input data.
+    ${ARRAY_FUNCTION_LIKE}
+
+        .. versionadded:: 1.20.0
 
     Returns
     -------
@@ -222,11 +275,24 @@ def asfortranarray(a, dtype=None):
     so it will not preserve 0-d arrays.  
 
     """
+    if like is not None:
+        return _asfortranarray_with_like(a, dtype=dtype, like=like)
+
     return array(a, dtype, copy=False, order='F', ndmin=1)
 
 
+_asfortranarray_with_like = array_function_dispatch(
+    _asarray_contiguous_fortran_dispatcher
+)(asfortranarray)
+
+
+def _require_dispatcher(a, dtype=None, requirements=None, *, like=None):
+    return (like,)
+
+
+@set_array_function_like_doc
 @set_module('numpy')
-def require(a, dtype=None, requirements=None):
+def require(a, dtype=None, requirements=None, *, like=None):
     """
     Return an ndarray of the provided type that satisfies requirements.
 
@@ -250,6 +316,9 @@ def require(a, dtype=None, requirements=None):
        * 'WRITEABLE' ('W')    - ensure a writable array
        * 'OWNDATA' ('O')      - ensure an array that owns its own data
        * 'ENSUREARRAY', ('E') - ensure a base array, instead of a subclass
+    ${ARRAY_FUNCTION_LIKE}
+
+        .. versionadded:: 1.20.0
 
     Returns
     -------
@@ -293,6 +362,14 @@ def require(a, dtype=None, requirements=None):
       UPDATEIFCOPY : False
 
     """
+    if like is not None:
+        return _require_with_like(
+            a,
+            dtype=dtype,
+            requirements=requirements,
+            like=like,
+        )
+
     possible_flags = {'C': 'C', 'C_CONTIGUOUS': 'C', 'CONTIGUOUS': 'C',
                       'F': 'F', 'F_CONTIGUOUS': 'F', 'FORTRAN': 'F',
                       'A': 'A', 'ALIGNED': 'A',
@@ -327,3 +404,8 @@ def require(a, dtype=None, requirements=None):
             arr = arr.copy(order)
             break
     return arr
+
+
+_require_with_like = array_function_dispatch(
+    _require_dispatcher
+)(require)