diff options
author | Peter Andreas Entschev <peter@entschev.com> | 2020-08-28 20:05:18 +0200 |
---|---|---|
committer | GitHub <noreply@github.com> | 2020-08-28 13:05:18 -0500 |
commit | 4cd6e4b336fbc68d88c0e9bc45a435ce7b721f1f (patch) | |
tree | 8823b0f72735aae4b836995760a8215fb8ef239f /numpy | |
parent | 32b3f82e4862d85c5aebba2b2b82a931038cd103 (diff) | |
download | numpy-4cd6e4b336fbc68d88c0e9bc45a435ce7b721f1f.tar.gz |
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>
Diffstat (limited to 'numpy')
-rw-r--r-- | numpy/__init__.pyi | 36 | ||||
-rw-r--r-- | numpy/core/_add_newdocs.py | 82 | ||||
-rw-r--r-- | numpy/core/_asarray.py | 94 | ||||
-rw-r--r-- | numpy/core/numeric.py | 76 | ||||
-rw-r--r-- | numpy/core/overrides.py | 21 | ||||
-rw-r--r-- | numpy/core/src/multiarray/arrayfunction_override.c | 170 | ||||
-rw-r--r-- | numpy/core/src/multiarray/arrayfunction_override.h | 4 | ||||
-rw-r--r-- | numpy/core/src/multiarray/multiarraymodule.c | 129 | ||||
-rw-r--r-- | numpy/core/src/multiarray/multiarraymodule.h | 2 | ||||
-rw-r--r-- | numpy/core/tests/test_overrides.py | 167 | ||||
-rw-r--r-- | numpy/lib/npyio.py | 65 | ||||
-rw-r--r-- | numpy/lib/twodim_base.py | 37 |
12 files changed, 772 insertions, 111 deletions
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi index c6cc94440..17f7dda67 100644 --- a/numpy/__init__.pyi +++ b/numpy/__init__.pyi @@ -524,15 +524,28 @@ def array( order: Optional[str] = ..., subok: bool = ..., ndmin: int = ..., + like: ArrayLike = ..., ) -> ndarray: ... def zeros( - shape: _ShapeLike, dtype: DtypeLike = ..., order: Optional[str] = ... + shape: _ShapeLike, + dtype: DtypeLike = ..., + order: Optional[str] = ..., + *, + like: ArrayLike = ..., ) -> ndarray: ... def ones( - shape: _ShapeLike, dtype: DtypeLike = ..., order: Optional[str] = ... + shape: _ShapeLike, + dtype: DtypeLike = ..., + order: Optional[str] = ..., + *, + like: ArrayLike = ..., ) -> ndarray: ... def empty( - shape: _ShapeLike, dtype: DtypeLike = ..., order: Optional[str] = ... + shape: _ShapeLike, + dtype: DtypeLike = ..., + order: Optional[str] = ..., + *, + like: ArrayLike = ..., ) -> ndarray: ... def zeros_like( a: ArrayLike, @@ -556,7 +569,12 @@ def empty_like( shape: Optional[_ShapeLike] = ..., ) -> ndarray: ... def full( - shape: _ShapeLike, fill_value: Any, dtype: DtypeLike = ..., order: str = ... + shape: _ShapeLike, + fill_value: Any, + dtype: DtypeLike = ..., + order: str = ..., + *, + like: ArrayLike = ..., ) -> ndarray: ... def full_like( a: ArrayLike, @@ -604,11 +622,17 @@ def cross( def indices( dimensions: Sequence[int], dtype: dtype = ..., sparse: bool = ... ) -> Union[ndarray, Tuple[ndarray, ...]]: ... -def fromfunction(function: Callable, shape: Tuple[int, int], **kwargs) -> Any: ... +def fromfunction( + function: Callable, + shape: Tuple[int, int], + *, + like: ArrayLike = ..., + **kwargs, +) -> Any: ... def isscalar(element: Any) -> bool: ... def binary_repr(num: int, width: Optional[int] = ...) -> str: ... def base_repr(number: int, base: int = ..., padding: int = ...) -> str: ... -def identity(n: int, dtype: DtypeLike = ...) -> ndarray: ... +def identity(n: int, dtype: DtypeLike = ..., *, like: ArrayLike = ...) -> ndarray: ... def allclose( a: ArrayLike, b: ArrayLike, diff --git a/numpy/core/_add_newdocs.py b/numpy/core/_add_newdocs.py index c3b4374f4..f11f008b2 100644 --- a/numpy/core/_add_newdocs.py +++ b/numpy/core/_add_newdocs.py @@ -12,6 +12,7 @@ NOTE: Many of the methods of ndarray have corresponding functions. from numpy.core import numerictypes as _numerictypes from numpy.core import dtype from numpy.core.function_base import add_newdoc +from numpy.core.overrides import array_function_like_doc ############################################################################### # @@ -786,7 +787,8 @@ add_newdoc('numpy.core', 'broadcast', ('reset', add_newdoc('numpy.core.multiarray', 'array', """ - array(object, dtype=None, *, copy=True, order='K', subok=False, ndmin=0) + array(object, dtype=None, *, copy=True, order='K', subok=False, ndmin=0, + like=None) Create an array. @@ -829,6 +831,9 @@ add_newdoc('numpy.core.multiarray', 'array', Specifies the minimum number of dimensions that the resulting array should have. Ones will be pre-pended to the shape as needed to meet this requirement. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -895,11 +900,14 @@ add_newdoc('numpy.core.multiarray', 'array', matrix([[1, 2], [3, 4]]) - """) + """.replace( + "${ARRAY_FUNCTION_LIKE}", + array_function_like_doc, + )) add_newdoc('numpy.core.multiarray', 'empty', """ - empty(shape, dtype=float, order='C') + empty(shape, dtype=float, order='C', *, like=None) Return a new array of given shape and type, without initializing entries. @@ -914,6 +922,9 @@ add_newdoc('numpy.core.multiarray', 'empty', Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -946,7 +957,10 @@ add_newdoc('numpy.core.multiarray', 'empty', array([[-1073741821, -1067949133], [ 496041986, 19249760]]) #uninitialized - """) + """.replace( + "${ARRAY_FUNCTION_LIKE}", + array_function_like_doc, + )) add_newdoc('numpy.core.multiarray', 'scalar', """ @@ -964,7 +978,7 @@ add_newdoc('numpy.core.multiarray', 'scalar', add_newdoc('numpy.core.multiarray', 'zeros', """ - zeros(shape, dtype=float, order='C') + zeros(shape, dtype=float, order='C', *, like=None) Return a new array of given shape and type, filled with zeros. @@ -979,6 +993,9 @@ add_newdoc('numpy.core.multiarray', 'zeros', Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -1013,7 +1030,10 @@ add_newdoc('numpy.core.multiarray', 'zeros', array([(0, 0), (0, 0)], dtype=[('x', '<i4'), ('y', '<i4')]) - """) + """.replace( + "${ARRAY_FUNCTION_LIKE}", + array_function_like_doc, + )) add_newdoc('numpy.core.multiarray', 'set_typeDict', """set_typeDict(dict) @@ -1025,7 +1045,7 @@ add_newdoc('numpy.core.multiarray', 'set_typeDict', add_newdoc('numpy.core.multiarray', 'fromstring', """ - fromstring(string, dtype=float, count=-1, sep='') + fromstring(string, dtype=float, count=-1, sep='', *, like=None) A new 1-D array initialized from text data in a string. @@ -1058,6 +1078,9 @@ add_newdoc('numpy.core.multiarray', 'fromstring', text, the binary mode of `fromstring` will first encode it into bytes using either utf-8 (python 3) or the default encoding (python 2), neither of which produce sane results. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -1081,7 +1104,10 @@ add_newdoc('numpy.core.multiarray', 'fromstring', >>> np.fromstring('1, 2', dtype=int, sep=',') array([1, 2]) - """) + """.replace( + "${ARRAY_FUNCTION_LIKE}", + array_function_like_doc, + )) add_newdoc('numpy.core.multiarray', 'compare_chararrays', """ @@ -1122,7 +1148,7 @@ add_newdoc('numpy.core.multiarray', 'compare_chararrays', add_newdoc('numpy.core.multiarray', 'fromiter', """ - fromiter(iterable, dtype, count=-1) + fromiter(iterable, dtype, count=-1, *, like=None) Create a new 1-dimensional array from an iterable object. @@ -1135,6 +1161,9 @@ add_newdoc('numpy.core.multiarray', 'fromiter', count : int, optional The number of items to read from *iterable*. The default is -1, which means all data is read. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -1152,11 +1181,14 @@ add_newdoc('numpy.core.multiarray', 'fromiter', >>> np.fromiter(iterable, float) array([ 0., 1., 4., 9., 16.]) - """) + """.replace( + "${ARRAY_FUNCTION_LIKE}", + array_function_like_doc, + )) add_newdoc('numpy.core.multiarray', 'fromfile', """ - fromfile(file, dtype=float, count=-1, sep='', offset=0) + fromfile(file, dtype=float, count=-1, sep='', offset=0, *, like=None) Construct an array from data in a text or binary file. @@ -1195,6 +1227,9 @@ add_newdoc('numpy.core.multiarray', 'fromfile', Only permitted for binary files. .. versionadded:: 1.17.0 + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 See also -------- @@ -1241,11 +1276,14 @@ add_newdoc('numpy.core.multiarray', 'fromfile', array([((10, 0), 98.25)], dtype=[('time', [('min', '<i8'), ('sec', '<i8')]), ('temp', '<f8')]) - """) + """.replace( + "${ARRAY_FUNCTION_LIKE}", + array_function_like_doc, + )) add_newdoc('numpy.core.multiarray', 'frombuffer', """ - frombuffer(buffer, dtype=float, count=-1, offset=0) + frombuffer(buffer, dtype=float, count=-1, offset=0, *, like=None) Interpret a buffer as a 1-dimensional array. @@ -1259,6 +1297,9 @@ add_newdoc('numpy.core.multiarray', 'frombuffer', Number of items to read. ``-1`` means all data in the buffer. offset : int, optional Start reading the buffer from this offset (in bytes); default: 0. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Notes ----- @@ -1283,7 +1324,10 @@ add_newdoc('numpy.core.multiarray', 'frombuffer', >>> np.frombuffer(b'\\x01\\x02\\x03\\x04\\x05', dtype=np.uint8, count=3) array([1, 2, 3], dtype=uint8) - """) + """.replace( + "${ARRAY_FUNCTION_LIKE}", + array_function_like_doc, + )) add_newdoc('numpy.core', 'fastCopyAndTranspose', """_fastCopyAndTranspose(a)""") @@ -1293,7 +1337,7 @@ add_newdoc('numpy.core.multiarray', 'correlate', add_newdoc('numpy.core.multiarray', 'arange', """ - arange([start,] stop[, step,], dtype=None) + arange([start,] stop[, step,], dtype=None, *, like=None) Return evenly spaced values within a given interval. @@ -1322,6 +1366,9 @@ add_newdoc('numpy.core.multiarray', 'arange', dtype : dtype The type of the output array. If `dtype` is not given, infer the data type from the other input arguments. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -1350,7 +1397,10 @@ add_newdoc('numpy.core.multiarray', 'arange', >>> np.arange(3,7,2) array([3, 5]) - """) + """.replace( + "${ARRAY_FUNCTION_LIKE}", + array_function_like_doc, + )) add_newdoc('numpy.core.multiarray', '_get_ndarray_c_version', """_get_ndarray_c_version() 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) diff --git a/numpy/core/numeric.py b/numpy/core/numeric.py index 84066dd30..a023bf0da 100644 --- a/numpy/core/numeric.py +++ b/numpy/core/numeric.py @@ -21,7 +21,7 @@ from .multiarray import ( from . import overrides from . import umath from . import shape_base -from .overrides import set_module +from .overrides import set_array_function_like_doc, set_module from .umath import (multiply, invert, sin, PINF, NAN) from . import numerictypes from .numerictypes import longlong, intc, int_, float_, complex_, bool_ @@ -141,8 +141,13 @@ def zeros_like(a, dtype=None, order='K', subok=True, shape=None): return res +def _ones_dispatcher(shape, dtype=None, order=None, *, like=None): + return(like,) + + +@set_array_function_like_doc @set_module('numpy') -def ones(shape, dtype=None, order='C'): +def ones(shape, dtype=None, order='C', *, like=None): """ Return a new array of given shape and type, filled with ones. @@ -157,6 +162,9 @@ def ones(shape, dtype=None, order='C'): Whether to store multi-dimensional data in row-major (C-style) or column-major (Fortran-style) order in memory. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -189,11 +197,19 @@ def ones(shape, dtype=None, order='C'): [1., 1.]]) """ + if like is not None: + return _ones_with_like(shape, dtype=dtype, order=order, like=like) + a = empty(shape, dtype, order) multiarray.copyto(a, 1, casting='unsafe') return a +_ones_with_like = array_function_dispatch( + _ones_dispatcher +)(ones) + + def _ones_like_dispatcher(a, dtype=None, order=None, subok=None, shape=None): return (a,) @@ -265,8 +281,13 @@ def ones_like(a, dtype=None, order='K', subok=True, shape=None): return res +def _full_dispatcher(shape, fill_value, dtype=None, order=None, *, like=None): + return(like,) + + +@set_array_function_like_doc @set_module('numpy') -def full(shape, fill_value, dtype=None, order='C'): +def full(shape, fill_value, dtype=None, order='C', *, like=None): """ Return a new array of given shape and type, filled with `fill_value`. @@ -282,6 +303,9 @@ def full(shape, fill_value, dtype=None, order='C'): order : {'C', 'F'}, optional Whether to store multidimensional data in C- or Fortran-contiguous (row- or column-wise) order in memory. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -309,6 +333,9 @@ def full(shape, fill_value, dtype=None, order='C'): [1, 2]]) """ + if like is not None: + return _full_with_like(shape, fill_value, dtype=dtype, order=order, like=like) + if dtype is None: fill_value = asarray(fill_value) dtype = fill_value.dtype @@ -317,6 +344,11 @@ def full(shape, fill_value, dtype=None, order='C'): return a +_full_with_like = array_function_dispatch( + _full_dispatcher +)(full) + + def _full_like_dispatcher(a, fill_value, dtype=None, order=None, subok=None, shape=None): return (a,) @@ -1754,8 +1786,13 @@ def indices(dimensions, dtype=int, sparse=False): return res +def _fromfunction_dispatcher(function, shape, *, dtype=None, like=None, **kwargs): + return (like,) + + +@set_array_function_like_doc @set_module('numpy') -def fromfunction(function, shape, *, dtype=float, **kwargs): +def fromfunction(function, shape, *, dtype=float, like=None, **kwargs): """ Construct an array by executing a function over each coordinate. @@ -1776,6 +1813,9 @@ def fromfunction(function, shape, *, dtype=float, **kwargs): dtype : data-type, optional Data-type of the coordinate arrays passed to `function`. By default, `dtype` is float. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -1806,10 +1846,18 @@ def fromfunction(function, shape, *, dtype=float, **kwargs): [2, 3, 4]]) """ + if like is not None: + return _fromfunction_with_like(function, shape, dtype=dtype, like=like, **kwargs) + args = indices(shape, dtype=dtype) return function(*args, **kwargs) +_fromfunction_with_like = array_function_dispatch( + _fromfunction_dispatcher +)(fromfunction) + + def _frombuffer(buf, dtype, shape, order): return frombuffer(buf, dtype=dtype).reshape(shape, order=order) @@ -2082,8 +2130,13 @@ def _maketup(descr, val): return tuple(res) +def _identity_dispatcher(n, dtype=None, *, like=None): + return (like,) + + +@set_array_function_like_doc @set_module('numpy') -def identity(n, dtype=None): +def identity(n, dtype=None, *, like=None): """ Return the identity array. @@ -2096,6 +2149,9 @@ def identity(n, dtype=None): Number of rows (and columns) in `n` x `n` output. dtype : data-type, optional Data-type of the output. Defaults to ``float``. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -2111,8 +2167,16 @@ def identity(n, dtype=None): [0., 0., 1.]]) """ + if like is not None: + return _identity_with_like(n, dtype=dtype, like=like) + from numpy import eye - return eye(n, dtype=dtype) + return eye(n, dtype=dtype, like=like) + + +_identity_with_like = array_function_dispatch( + _identity_dispatcher +)(identity) def _allclose_dispatcher(a, b, rtol=None, atol=None, equal_nan=None): diff --git a/numpy/core/overrides.py b/numpy/core/overrides.py index 816b11293..c2b5fb7fa 100644 --- a/numpy/core/overrides.py +++ b/numpy/core/overrides.py @@ -12,6 +12,27 @@ from numpy.compat._inspect import getargspec ARRAY_FUNCTION_ENABLED = bool( int(os.environ.get('NUMPY_EXPERIMENTAL_ARRAY_FUNCTION', 1))) +array_function_like_doc = ( + """like : array_like + Reference object to allow the creation of arrays which are not + NumPy arrays. If an array-like passed in as ``like`` supports + the ``__array_function__`` protocol, the result will be defined + by it. In this case, it ensures the creation of an array object + compatible with that passed in via this argument. + + .. note:: + The ``like`` keyword is an experimental feature pending on + acceptance of :ref:`NEP 35 <NEP35>`.""" +) + +def set_array_function_like_doc(public_api): + if public_api.__doc__ is not None: + public_api.__doc__ = public_api.__doc__.replace( + "${ARRAY_FUNCTION_LIKE}", + array_function_like_doc, + ) + return public_api + add_docstring( implement_array_function, diff --git a/numpy/core/src/multiarray/arrayfunction_override.c b/numpy/core/src/multiarray/arrayfunction_override.c index 9ea8efdd9..613fe6b3f 100644 --- a/numpy/core/src/multiarray/arrayfunction_override.c +++ b/numpy/core/src/multiarray/arrayfunction_override.c @@ -26,7 +26,6 @@ static PyObject * get_array_function(PyObject *obj) { static PyObject *ndarray_array_function = NULL; - PyObject *array_function; if (ndarray_array_function == NULL) { ndarray_array_function = get_ndarray_array_function(); @@ -38,7 +37,7 @@ get_array_function(PyObject *obj) return ndarray_array_function; } - array_function = PyArray_LookupSpecial(obj, "__array_function__"); + PyObject *array_function = PyArray_LookupSpecial(obj, "__array_function__"); if (array_function == NULL && PyErr_Occurred()) { PyErr_Clear(); /* TODO[gh-14801]: propagate crashes during attribute access? */ } @@ -53,9 +52,7 @@ get_array_function(PyObject *obj) static void pyobject_array_insert(PyObject **array, int length, int index, PyObject *item) { - int j; - - for (j = length; j > index; j--) { + for (int j = length; j > index; j--) { array[j] = array[j - 1]; } array[index] = item; @@ -74,18 +71,16 @@ get_implementing_args_and_methods(PyObject *relevant_args, PyObject **methods) { int num_implementing_args = 0; - Py_ssize_t i; - int j; PyObject **items = PySequence_Fast_ITEMS(relevant_args); Py_ssize_t length = PySequence_Fast_GET_SIZE(relevant_args); - for (i = 0; i < length; i++) { + for (Py_ssize_t i = 0; i < length; i++) { int new_class = 1; PyObject *argument = items[i]; /* Have we seen this type before? */ - for (j = 0; j < num_implementing_args; j++) { + for (int j = 0; j < num_implementing_args; j++) { if (Py_TYPE(argument) == Py_TYPE(implementing_args[j])) { new_class = 0; break; @@ -109,7 +104,7 @@ get_implementing_args_and_methods(PyObject *relevant_args, /* "subclasses before superclasses, otherwise left to right" */ arg_index = num_implementing_args; - for (j = 0; j < num_implementing_args; j++) { + for (int j = 0; j < num_implementing_args; j++) { PyObject *other_type; other_type = (PyObject *)Py_TYPE(implementing_args[j]); if (PyObject_IsInstance(argument, other_type)) { @@ -129,7 +124,7 @@ get_implementing_args_and_methods(PyObject *relevant_args, return num_implementing_args; fail: - for (j = 0; j < num_implementing_args; j++) { + for (int j = 0; j < num_implementing_args; j++) { Py_DECREF(implementing_args[j]); Py_DECREF(methods[j]); } @@ -161,13 +156,10 @@ NPY_NO_EXPORT PyObject * array_function_method_impl(PyObject *func, PyObject *types, PyObject *args, PyObject *kwargs) { - Py_ssize_t j; - PyObject *implementation, *result; - PyObject **items = PySequence_Fast_ITEMS(types); Py_ssize_t length = PySequence_Fast_GET_SIZE(types); - for (j = 0; j < length; j++) { + for (Py_ssize_t j = 0; j < length; j++) { int is_subclass = PyObject_IsSubclass( items[j], (PyObject *)&PyArray_Type); if (is_subclass == -1) { @@ -179,11 +171,11 @@ array_function_method_impl(PyObject *func, PyObject *types, PyObject *args, } } - implementation = PyObject_GetAttr(func, npy_ma_str_implementation); + PyObject *implementation = PyObject_GetAttr(func, npy_ma_str_implementation); if (implementation == NULL) { return NULL; } - result = PyObject_Call(implementation, args, kwargs); + PyObject *result = PyObject_Call(implementation, args, kwargs); Py_DECREF(implementation); return result; } @@ -208,32 +200,32 @@ call_array_function(PyObject* argument, PyObject* method, } -/* - * Implements the __array_function__ protocol for a function, as described in - * in NEP-18. See numpy.core.overrides for a full docstring. +/** + * Internal handler for the array-function dispatching. The helper returns + * either the result, or NotImplemented (as a borrowed reference). + * + * @param public_api The public API symbol used for dispatching + * @param relevant_args Arguments which may implement __array_function__ + * @param args Original arguments + * @param kwargs Original keyword arguments + * + * @returns The result of the dispatched version, or a borrowed reference + * to NotImplemented to indicate the default implementation should + * be used. */ NPY_NO_EXPORT PyObject * -array_implement_array_function( - PyObject *NPY_UNUSED(dummy), PyObject *positional_args) +array_implement_array_function_internal( + PyObject *public_api, PyObject *relevant_args, + PyObject *args, PyObject *kwargs) { - PyObject *implementation, *public_api, *relevant_args, *args, *kwargs; - - PyObject *types = NULL; PyObject *implementing_args[NPY_MAXARGS]; PyObject *array_function_methods[NPY_MAXARGS]; + PyObject *types = NULL; - int j, any_overrides; - int num_implementing_args = 0; PyObject *result = NULL; static PyObject *errmsg_formatter = NULL; - if (!PyArg_UnpackTuple( - positional_args, "implement_array_function", 5, 5, - &implementation, &public_api, &relevant_args, &args, &kwargs)) { - return NULL; - } - relevant_args = PySequence_Fast( relevant_args, "dispatcher for __array_function__ did not return an iterable"); @@ -242,7 +234,7 @@ array_implement_array_function( } /* Collect __array_function__ implementations */ - num_implementing_args = get_implementing_args_and_methods( + int num_implementing_args = get_implementing_args_and_methods( relevant_args, implementing_args, array_function_methods); if (num_implementing_args == -1) { goto cleanup; @@ -254,15 +246,19 @@ array_implement_array_function( * arguments implement __array_function__ at all (e.g., if they are all * built-in types). */ - any_overrides = 0; - for (j = 0; j < num_implementing_args; j++) { + int any_overrides = 0; + for (int j = 0; j < num_implementing_args; j++) { if (!is_default_array_function(array_function_methods[j])) { any_overrides = 1; break; } } if (!any_overrides) { - result = PyObject_Call(implementation, args, kwargs); + /* + * When the default implementation should be called, return + * `Py_NotImplemented` to indicate this. + */ + result = Py_NotImplemented; goto cleanup; } @@ -275,14 +271,14 @@ array_implement_array_function( if (types == NULL) { goto cleanup; } - for (j = 0; j < num_implementing_args; j++) { + for (int j = 0; j < num_implementing_args; j++) { PyObject *arg_type = (PyObject *)Py_TYPE(implementing_args[j]); Py_INCREF(arg_type); PyTuple_SET_ITEM(types, j, arg_type); } /* Call __array_function__ methods */ - for (j = 0; j < num_implementing_args; j++) { + for (int j = 0; j < num_implementing_args; j++) { PyObject *argument = implementing_args[j]; PyObject *method = array_function_methods[j]; @@ -319,7 +315,7 @@ array_implement_array_function( } cleanup: - for (j = 0; j < num_implementing_args; j++) { + for (int j = 0; j < num_implementing_args; j++) { Py_DECREF(implementing_args[j]); Py_DECREF(array_function_methods[j]); } @@ -330,6 +326,92 @@ cleanup: /* + * Implements the __array_function__ protocol for a Python function, as described in + * in NEP-18. See numpy.core.overrides for a full docstring. + */ +NPY_NO_EXPORT PyObject * +array_implement_array_function( + PyObject *NPY_UNUSED(dummy), PyObject *positional_args) +{ + PyObject *implementation, *public_api, *relevant_args, *args, *kwargs; + + if (!PyArg_UnpackTuple( + positional_args, "implement_array_function", 5, 5, + &implementation, &public_api, &relevant_args, &args, &kwargs)) { + return NULL; + } + + /* Remove `like=` kwarg, which is NumPy-exclusive and thus not present + * in downstream libraries. + */ + if (kwargs != NULL && PyDict_Contains(kwargs, npy_ma_str_like)) { + PyDict_DelItem(kwargs, npy_ma_str_like); + } + + PyObject *res = array_implement_array_function_internal( + public_api, relevant_args, args, kwargs); + + if (res == Py_NotImplemented) { + return PyObject_Call(implementation, args, kwargs); + } + return res; +} + + +/* + * Implements the __array_function__ protocol for C array creation functions + * only. Added as an extension to NEP-18 in an effort to bring NEP-35 to + * life with minimal dispatch overhead. + */ +NPY_NO_EXPORT PyObject * +array_implement_c_array_function_creation( + const char *function_name, PyObject *args, PyObject *kwargs) +{ + if (kwargs == NULL) { + return Py_NotImplemented; + } + + /* Remove `like=` kwarg, which is NumPy-exclusive and thus not present + * in downstream libraries. If that key isn't present, return NULL and + * let originating call to continue. + */ + if (!PyDict_Contains(kwargs, npy_ma_str_like)) { + return Py_NotImplemented; + } + + PyObject *relevant_args = PyTuple_Pack(1, + PyDict_GetItem(kwargs, npy_ma_str_like)); + if (relevant_args == NULL) { + return NULL; + } + PyDict_DelItem(kwargs, npy_ma_str_like); + + PyObject *numpy_module = PyImport_Import(npy_ma_str_numpy); + if (numpy_module == NULL) { + return NULL; + } + + PyObject *public_api = PyObject_GetAttrString(numpy_module, function_name); + Py_DECREF(numpy_module); + if (public_api == NULL) { + return NULL; + } + if (!PyCallable_Check(public_api)) { + Py_DECREF(public_api); + return PyErr_Format(PyExc_RuntimeError, + "numpy.%s is not callable.", + function_name); + } + + PyObject* result = array_implement_array_function_internal( + public_api, relevant_args, args, kwargs); + + Py_DECREF(public_api); + return result; +} + + +/* * Python wrapper for get_implementing_args_and_methods, for testing purposes. */ NPY_NO_EXPORT PyObject * @@ -337,8 +419,6 @@ array__get_implementing_args( PyObject *NPY_UNUSED(dummy), PyObject *positional_args) { PyObject *relevant_args; - int j; - int num_implementing_args = 0; PyObject *implementing_args[NPY_MAXARGS]; PyObject *array_function_methods[NPY_MAXARGS]; PyObject *result = NULL; @@ -355,7 +435,7 @@ array__get_implementing_args( return NULL; } - num_implementing_args = get_implementing_args_and_methods( + int num_implementing_args = get_implementing_args_and_methods( relevant_args, implementing_args, array_function_methods); if (num_implementing_args == -1) { goto cleanup; @@ -366,14 +446,14 @@ array__get_implementing_args( if (result == NULL) { goto cleanup; } - for (j = 0; j < num_implementing_args; j++) { + for (int j = 0; j < num_implementing_args; j++) { PyObject *argument = implementing_args[j]; Py_INCREF(argument); PyList_SET_ITEM(result, j, argument); } cleanup: - for (j = 0; j < num_implementing_args; j++) { + for (int j = 0; j < num_implementing_args; j++) { Py_DECREF(implementing_args[j]); Py_DECREF(array_function_methods[j]); } diff --git a/numpy/core/src/multiarray/arrayfunction_override.h b/numpy/core/src/multiarray/arrayfunction_override.h index 0d224e2b6..fdcf1746d 100644 --- a/numpy/core/src/multiarray/arrayfunction_override.h +++ b/numpy/core/src/multiarray/arrayfunction_override.h @@ -10,6 +10,10 @@ array__get_implementing_args( PyObject *NPY_UNUSED(dummy), PyObject *positional_args); NPY_NO_EXPORT PyObject * +array_implement_c_array_function_creation( + const char *function_name, PyObject *args, PyObject *kwargs); + +NPY_NO_EXPORT PyObject * array_function_method_impl(PyObject *func, PyObject *types, PyObject *args, PyObject *kwargs); diff --git a/numpy/core/src/multiarray/multiarraymodule.c b/numpy/core/src/multiarray/multiarraymodule.c index c79d9a845..bc367b78d 100644 --- a/numpy/core/src/multiarray/multiarraymodule.c +++ b/numpy/core/src/multiarray/multiarraymodule.c @@ -1582,13 +1582,16 @@ _array_fromobject(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kws) npy_bool subok = NPY_FALSE; npy_bool copy = NPY_TRUE; int ndmin = 0, nd; + PyObject* like; PyArray_Descr *type = NULL; PyArray_Descr *oldtype = NULL; NPY_ORDER order = NPY_KEEPORDER; int flags = 0; - static char *kwd[]= {"object", "dtype", "copy", "order", "subok", - "ndmin", NULL}; + PyObject* array_function_result = NULL; + + static char *kwd[] = {"object", "dtype", "copy", "order", "subok", + "ndmin", "like", NULL}; if (PyTuple_GET_SIZE(args) > 2) { PyErr_Format(PyExc_TypeError, @@ -1597,6 +1600,12 @@ _array_fromobject(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kws) return NULL; } + array_function_result = array_implement_c_array_function_creation( + "array", args, kws); + if (array_function_result != Py_NotImplemented) { + return array_function_result; + } + /* super-fast path for ndarray argument calls */ if (PyTuple_GET_SIZE(args) == 0) { goto full_path; @@ -1674,13 +1683,14 @@ _array_fromobject(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kws) } full_path: - if (!PyArg_ParseTupleAndKeywords(args, kws, "O|O&O&O&O&i:array", kwd, + if (!PyArg_ParseTupleAndKeywords(args, kws, "O|O&O&O&O&i$O:array", kwd, &op, PyArray_DescrConverter2, &type, PyArray_BoolConverter, ©, PyArray_OrderConverter, &order, PyArray_BoolConverter, &subok, - &ndmin)) { + &ndmin, + &like)) { goto clean_type; } @@ -1817,20 +1827,29 @@ static PyObject * array_empty(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kwds) { - static char *kwlist[] = {"shape", "dtype", "order", NULL}; + static char *kwlist[] = {"shape", "dtype", "order", "like", NULL}; PyArray_Descr *typecode = NULL; PyArray_Dims shape = {NULL, 0}; NPY_ORDER order = NPY_CORDER; + PyObject *like = NULL; npy_bool is_f_order; + PyObject *array_function_result = NULL; PyArrayObject *ret = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|O&O&:empty", kwlist, + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|O&O&$O:empty", kwlist, PyArray_IntpConverter, &shape, PyArray_DescrConverter, &typecode, - PyArray_OrderConverter, &order)) { + PyArray_OrderConverter, &order, + &like)) { goto fail; } + array_function_result = array_implement_c_array_function_creation( + "empty", args, kwds); + if (array_function_result != Py_NotImplemented) { + return array_function_result; + } + switch (order) { case NPY_CORDER: is_f_order = NPY_FALSE; @@ -1984,20 +2003,29 @@ array_scalar(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kwds) static PyObject * array_zeros(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kwds) { - static char *kwlist[] = {"shape", "dtype", "order", NULL}; + static char *kwlist[] = {"shape", "dtype", "order", "like", NULL}; PyArray_Descr *typecode = NULL; PyArray_Dims shape = {NULL, 0}; NPY_ORDER order = NPY_CORDER; + PyObject *like = NULL; npy_bool is_f_order = NPY_FALSE; + PyObject *array_function_result = NULL; PyArrayObject *ret = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|O&O&:zeros", kwlist, + if (!PyArg_ParseTupleAndKeywords(args, kwds, "O&|O&O&$O:zeros", kwlist, PyArray_IntpConverter, &shape, PyArray_DescrConverter, &typecode, - PyArray_OrderConverter, &order)) { + PyArray_OrderConverter, &order, + &like)) { goto fail; } + array_function_result = array_implement_c_array_function_creation( + "zeros", args, kwds); + if (array_function_result != Py_NotImplemented) { + return array_function_result; + } + switch (order) { case NPY_CORDER: is_f_order = NPY_FALSE; @@ -2050,16 +2078,24 @@ array_fromstring(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *keywds Py_ssize_t nin = -1; char *sep = NULL; Py_ssize_t s; - static char *kwlist[] = {"string", "dtype", "count", "sep", NULL}; + static char *kwlist[] = {"string", "dtype", "count", "sep", "like", NULL}; + PyObject *like = NULL; PyArray_Descr *descr = NULL; + PyObject *array_function_result = NULL; if (!PyArg_ParseTupleAndKeywords(args, keywds, - "s#|O&" NPY_SSIZE_T_PYFMT "s:fromstring", kwlist, - &data, &s, PyArray_DescrConverter, &descr, &nin, &sep)) { + "s#|O&" NPY_SSIZE_T_PYFMT "s$O:fromstring", kwlist, + &data, &s, PyArray_DescrConverter, &descr, &nin, &sep, &like)) { Py_XDECREF(descr); return NULL; } + array_function_result = array_implement_c_array_function_creation( + "fromstring", args, keywds); + if (array_function_result != Py_NotImplemented) { + return array_function_result; + } + /* binary mode, condition copied from PyArray_FromString */ if (sep == NULL || strlen(sep) == 0) { /* Numpy 1.14, 2017-10-19 */ @@ -2082,19 +2118,27 @@ array_fromfile(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *keywds) PyObject *err_type = NULL, *err_value = NULL, *err_traceback = NULL; char *sep = ""; Py_ssize_t nin = -1; - static char *kwlist[] = {"file", "dtype", "count", "sep", "offset", NULL}; + static char *kwlist[] = {"file", "dtype", "count", "sep", "offset", "like", NULL}; + PyObject *like = NULL; PyArray_Descr *type = NULL; + PyObject *array_function_result = NULL; int own; npy_off_t orig_pos = 0, offset = 0; FILE *fp; if (!PyArg_ParseTupleAndKeywords(args, keywds, - "O|O&" NPY_SSIZE_T_PYFMT "s" NPY_OFF_T_PYFMT ":fromfile", kwlist, - &file, PyArray_DescrConverter, &type, &nin, &sep, &offset)) { + "O|O&" NPY_SSIZE_T_PYFMT "s" NPY_OFF_T_PYFMT "$O:fromfile", kwlist, + &file, PyArray_DescrConverter, &type, &nin, &sep, &offset, &like)) { Py_XDECREF(type); return NULL; } + array_function_result = array_implement_c_array_function_creation( + "fromfile", args, keywds); + if (array_function_result != Py_NotImplemented) { + return array_function_result; + } + file = NpyPath_PathlikeToFspath(file); if (file == NULL) { return NULL; @@ -2161,15 +2205,24 @@ array_fromiter(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *keywds) { PyObject *iter; Py_ssize_t nin = -1; - static char *kwlist[] = {"iter", "dtype", "count", NULL}; + static char *kwlist[] = {"iter", "dtype", "count", "like", NULL}; + PyObject *like = NULL; PyArray_Descr *descr = NULL; + PyObject *array_function_result = NULL; if (!PyArg_ParseTupleAndKeywords(args, keywds, - "OO&|" NPY_SSIZE_T_PYFMT ":fromiter", kwlist, - &iter, PyArray_DescrConverter, &descr, &nin)) { + "OO&|" NPY_SSIZE_T_PYFMT "$O:fromiter", kwlist, + &iter, PyArray_DescrConverter, &descr, &nin, &like)) { Py_XDECREF(descr); return NULL; } + + array_function_result = array_implement_c_array_function_creation( + "fromiter", args, keywds); + if (array_function_result != Py_NotImplemented) { + return array_function_result; + } + return PyArray_FromIter(iter, descr, (npy_intp)nin); } @@ -2178,15 +2231,24 @@ array_frombuffer(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *keywds { PyObject *obj = NULL; Py_ssize_t nin = -1, offset = 0; - static char *kwlist[] = {"buffer", "dtype", "count", "offset", NULL}; + static char *kwlist[] = {"buffer", "dtype", "count", "offset", "like", NULL}; + PyObject *like = NULL; PyArray_Descr *type = NULL; + PyObject *array_function_result = NULL; if (!PyArg_ParseTupleAndKeywords(args, keywds, - "O|O&" NPY_SSIZE_T_PYFMT NPY_SSIZE_T_PYFMT ":frombuffer", kwlist, - &obj, PyArray_DescrConverter, &type, &nin, &offset)) { + "O|O&" NPY_SSIZE_T_PYFMT NPY_SSIZE_T_PYFMT "$O:frombuffer", kwlist, + &obj, PyArray_DescrConverter, &type, &nin, &offset, &like)) { Py_XDECREF(type); return NULL; } + + array_function_result = array_implement_c_array_function_creation( + "frombuffer", args, keywds); + if (array_function_result != Py_NotImplemented) { + return array_function_result; + } + if (type == NULL) { type = PyArray_DescrFromType(NPY_DEFAULT_TYPE); } @@ -2766,17 +2828,27 @@ array_correlate2(PyObject *NPY_UNUSED(dummy), PyObject *args, PyObject *kwds) static PyObject * array_arange(PyObject *NPY_UNUSED(ignored), PyObject *args, PyObject *kws) { PyObject *o_start = NULL, *o_stop = NULL, *o_step = NULL, *range=NULL; - static char *kwd[]= {"start", "stop", "step", "dtype", NULL}; + PyObject *like = NULL; + PyObject *array_function_result = NULL; + static char *kwd[] = {"start", "stop", "step", "dtype", "like", NULL}; PyArray_Descr *typecode = NULL; - if (!PyArg_ParseTupleAndKeywords(args, kws, "O|OOO&:arange", kwd, + if (!PyArg_ParseTupleAndKeywords(args, kws, "O|OOO&$O:arange", kwd, &o_start, &o_stop, &o_step, - PyArray_DescrConverter2, &typecode)) { + PyArray_DescrConverter2, &typecode, + &like)) { Py_XDECREF(typecode); return NULL; } + + array_function_result = array_implement_c_array_function_creation( + "arange", args, kws); + if (array_function_result != Py_NotImplemented) { + return array_function_result; + } + range = PyArray_ArangeObj(o_start, o_stop, o_step, typecode); Py_XDECREF(typecode); @@ -4331,6 +4403,8 @@ NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_dtype = NULL; NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_ndmin = NULL; NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_axis1 = NULL; NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_axis2 = NULL; +NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_like = NULL; +NPY_VISIBILITY_HIDDEN PyObject * npy_ma_str_numpy = NULL; static int intern_strings(void) @@ -4347,12 +4421,15 @@ intern_strings(void) npy_ma_str_ndmin = PyUnicode_InternFromString("ndmin"); npy_ma_str_axis1 = PyUnicode_InternFromString("axis1"); npy_ma_str_axis2 = PyUnicode_InternFromString("axis2"); + npy_ma_str_like = PyUnicode_InternFromString("like"); + npy_ma_str_numpy = PyUnicode_InternFromString("numpy"); return npy_ma_str_array && npy_ma_str_array_prepare && npy_ma_str_array_wrap && npy_ma_str_array_finalize && npy_ma_str_ufunc && npy_ma_str_implementation && npy_ma_str_order && npy_ma_str_copy && npy_ma_str_dtype && - npy_ma_str_ndmin && npy_ma_str_axis1 && npy_ma_str_axis2; + npy_ma_str_ndmin && npy_ma_str_axis1 && npy_ma_str_axis2 && + npy_ma_str_like && npy_ma_str_numpy; } static struct PyModuleDef moduledef = { diff --git a/numpy/core/src/multiarray/multiarraymodule.h b/numpy/core/src/multiarray/multiarraymodule.h index dd437e091..d3ee3337c 100644 --- a/numpy/core/src/multiarray/multiarraymodule.h +++ b/numpy/core/src/multiarray/multiarraymodule.h @@ -13,5 +13,7 @@ NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_dtype; NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_ndmin; NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_axis1; NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_axis2; +NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_like; +NPY_VISIBILITY_HIDDEN extern PyObject * npy_ma_str_numpy; #endif diff --git a/numpy/core/tests/test_overrides.py b/numpy/core/tests/test_overrides.py index 7e73d8c03..42600a12b 100644 --- a/numpy/core/tests/test_overrides.py +++ b/numpy/core/tests/test_overrides.py @@ -1,5 +1,7 @@ import inspect import sys +import tempfile +from io import StringIO from unittest import mock import numpy as np @@ -425,3 +427,168 @@ class TestNumPyFunctions: # note: the internal implementation of np.sum() calls the .sum() method array = np.array(1).view(MyArray) assert_equal(np.sum(array), 'summed') + + +class TestArrayLike: + + class MyArray(): + + def __init__(self, function=None): + self.function = function + + def __array_function__(self, func, types, args, kwargs): + try: + my_func = getattr(TestArrayLike.MyArray, func.__name__) + except AttributeError: + return NotImplemented + return my_func(*args, **kwargs) + + class MyNoArrayFunctionArray(): + + def __init__(self, function=None): + self.function = function + + def add_method(name, arr_class, enable_value_error=False): + def _definition(*args, **kwargs): + # Check that `like=` isn't propagated downstream + assert 'like' not in kwargs + + if enable_value_error and 'value_error' in kwargs: + raise ValueError + + return arr_class(getattr(arr_class, name)) + setattr(arr_class, name, _definition) + + def func_args(*args, **kwargs): + return args, kwargs + + @requires_array_function + def test_array_like_not_implemented(self): + TestArrayLike.add_method('array', TestArrayLike.MyArray) + + ref = TestArrayLike.MyArray.array() + + with assert_raises_regex(TypeError, 'no implementation found'): + array_like = np.asarray(1, like=ref) + + _array_tests = [ + ('array', *func_args((1,))), + ('asarray', *func_args((1,))), + ('asanyarray', *func_args((1,))), + ('ascontiguousarray', *func_args((2, 3))), + ('asfortranarray', *func_args((2, 3))), + ('require', *func_args((np.arange(6).reshape(2, 3),), + requirements=['A', 'F'])), + ('empty', *func_args((1,))), + ('full', *func_args((1,), 2)), + ('ones', *func_args((1,))), + ('zeros', *func_args((1,))), + ('arange', *func_args(3)), + ('frombuffer', *func_args(b'\x00' * 8, dtype=int)), + ('fromiter', *func_args(range(3), dtype=int)), + ('fromstring', *func_args('1,2', dtype=int, sep=',')), + ('loadtxt', *func_args(lambda: StringIO('0 1\n2 3'))), + ('genfromtxt', *func_args(lambda: StringIO(u'1,2.1'), + dtype=[('int', 'i8'), ('float', 'f8')], + delimiter=',')), + ] + + @pytest.mark.parametrize('function, args, kwargs', _array_tests) + @pytest.mark.parametrize('numpy_ref', [True, False]) + @requires_array_function + def test_array_like(self, function, args, kwargs, numpy_ref): + TestArrayLike.add_method('array', TestArrayLike.MyArray) + TestArrayLike.add_method(function, TestArrayLike.MyArray) + np_func = getattr(np, function) + my_func = getattr(TestArrayLike.MyArray, function) + + if numpy_ref is True: + ref = np.array(1) + else: + ref = TestArrayLike.MyArray.array() + + like_args = tuple(a() if callable(a) else a for a in args) + array_like = np_func(*like_args, **kwargs, like=ref) + + if numpy_ref is True: + assert type(array_like) is np.ndarray + + np_args = tuple(a() if callable(a) else a for a in args) + np_arr = np_func(*np_args, **kwargs) + + # Special-case np.empty to ensure values match + if function == "empty": + np_arr.fill(1) + array_like.fill(1) + + assert_equal(array_like, np_arr) + else: + assert type(array_like) is TestArrayLike.MyArray + assert array_like.function is my_func + + @pytest.mark.parametrize('function, args, kwargs', _array_tests) + @pytest.mark.parametrize('numpy_ref', [True, False]) + @requires_array_function + def test_no_array_function_like(self, function, args, kwargs, numpy_ref): + TestArrayLike.add_method('array', TestArrayLike.MyNoArrayFunctionArray) + TestArrayLike.add_method(function, TestArrayLike.MyNoArrayFunctionArray) + np_func = getattr(np, function) + my_func = getattr(TestArrayLike.MyNoArrayFunctionArray, function) + + if numpy_ref is True: + ref = np.array(1) + else: + ref = TestArrayLike.MyNoArrayFunctionArray.array() + + like_args = tuple(a() if callable(a) else a for a in args) + array_like = np_func(*like_args, **kwargs, like=ref) + + assert type(array_like) is np.ndarray + if numpy_ref is True: + np_args = tuple(a() if callable(a) else a for a in args) + np_arr = np_func(*np_args, **kwargs) + + # Special-case np.empty to ensure values match + if function == "empty": + np_arr.fill(1) + array_like.fill(1) + + assert_equal(array_like, np_arr) + + @pytest.mark.parametrize('numpy_ref', [True, False]) + def test_array_like_fromfile(self, numpy_ref): + TestArrayLike.add_method('array', TestArrayLike.MyArray) + TestArrayLike.add_method("fromfile", TestArrayLike.MyArray) + + if numpy_ref is True: + ref = np.array(1) + else: + ref = TestArrayLike.MyArray.array() + + data = np.random.random(5) + + fname = tempfile.mkstemp()[1] + data.tofile(fname) + + array_like = np.fromfile(fname, like=ref) + if numpy_ref is True: + assert type(array_like) is np.ndarray + np_res = np.fromfile(fname, like=ref) + assert_equal(np_res, data) + assert_equal(array_like, np_res) + else: + assert type(array_like) is TestArrayLike.MyArray + assert array_like.function is TestArrayLike.MyArray.fromfile + + @requires_array_function + def test_exception_handling(self): + TestArrayLike.add_method( + 'array', + TestArrayLike.MyArray, + enable_value_error=True, + ) + + ref = TestArrayLike.MyArray.array() + + with assert_raises(ValueError): + np.array(1, value_error=True, like=ref) diff --git a/numpy/lib/npyio.py b/numpy/lib/npyio.py index 58affc2fc..cc3465cc6 100644 --- a/numpy/lib/npyio.py +++ b/numpy/lib/npyio.py @@ -14,7 +14,7 @@ from . import format from ._datasource import DataSource from numpy.core import overrides from numpy.core.multiarray import packbits, unpackbits -from numpy.core.overrides import set_module +from numpy.core.overrides import set_array_function_like_doc, set_module from numpy.core._internal import recursive from ._iotools import ( LineSplitter, NameValidator, StringConverter, ConverterError, @@ -790,10 +790,17 @@ def _getconv(dtype): _loadtxt_chunksize = 50000 +def _loadtxt_dispatcher(fname, dtype=None, comments=None, delimiter=None, + converters=None, skiprows=None, usecols=None, unpack=None, + ndmin=None, encoding=None, max_rows=None, *, like=None): + return (like,) + + +@set_array_function_like_doc @set_module('numpy') def loadtxt(fname, dtype=float, comments='#', delimiter=None, converters=None, skiprows=0, usecols=None, unpack=False, - ndmin=0, encoding='bytes', max_rows=None): + ndmin=0, encoding='bytes', max_rows=None, *, like=None): r""" Load data from a text file. @@ -860,6 +867,9 @@ def loadtxt(fname, dtype=float, comments='#', delimiter=None, is to read all the lines. .. versionadded:: 1.16.0 + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -917,6 +927,14 @@ def loadtxt(fname, dtype=float, comments='#', delimiter=None, [-17.57, 63.94]]) """ + if like is not None: + return _loadtxt_with_like( + fname, dtype=dtype, comments=comments, delimiter=delimiter, + converters=converters, skiprows=skiprows, usecols=usecols, + unpack=unpack, ndmin=ndmin, encoding=encoding, + max_rows=max_rows, like=like + ) + # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - # Nested functions used by loadtxt. # - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - @@ -1201,6 +1219,11 @@ def loadtxt(fname, dtype=float, comments='#', delimiter=None, return X +_loadtxt_with_like = array_function_dispatch( + _loadtxt_dispatcher +)(loadtxt) + + def _savetxt_dispatcher(fname, X, fmt=None, delimiter=None, newline=None, header=None, footer=None, comments=None, encoding=None): @@ -1554,6 +1577,18 @@ def fromregex(file, regexp, dtype, encoding=None): #####-------------------------------------------------------------------------- +def _genfromtxt_dispatcher(fname, dtype=None, comments=None, delimiter=None, + skip_header=None, skip_footer=None, converters=None, + missing_values=None, filling_values=None, usecols=None, + names=None, excludelist=None, deletechars=None, + replace_space=None, autostrip=None, case_sensitive=None, + defaultfmt=None, unpack=None, usemask=None, loose=None, + invalid_raise=None, max_rows=None, encoding=None, *, + like=None): + return (like,) + + +@set_array_function_like_doc @set_module('numpy') def genfromtxt(fname, dtype=float, comments='#', delimiter=None, skip_header=0, skip_footer=0, converters=None, @@ -1562,7 +1597,8 @@ def genfromtxt(fname, dtype=float, comments='#', delimiter=None, deletechars=''.join(sorted(NameValidator.defaultdeletechars)), replace_space='_', autostrip=False, case_sensitive=True, defaultfmt="f%i", unpack=None, usemask=False, loose=True, - invalid_raise=True, max_rows=None, encoding='bytes'): + invalid_raise=True, max_rows=None, encoding='bytes', *, + like=None): """ Load data from a text file, with missing values handled as specified. @@ -1659,6 +1695,9 @@ def genfromtxt(fname, dtype=float, comments='#', delimiter=None, to None the system default is used. The default value is 'bytes'. .. versionadded:: 1.14.0 + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -1737,6 +1776,21 @@ def genfromtxt(fname, dtype=float, comments='#', delimiter=None, dtype=[('f0', 'S12'), ('f1', 'S12')]) """ + + if like is not None: + return _genfromtxt_with_like( + fname, dtype=dtype, comments=comments, delimiter=delimiter, + skip_header=skip_header, skip_footer=skip_footer, + converters=converters, missing_values=missing_values, + filling_values=filling_values, usecols=usecols, names=names, + excludelist=excludelist, deletechars=deletechars, + replace_space=replace_space, autostrip=autostrip, + case_sensitive=case_sensitive, defaultfmt=defaultfmt, + unpack=unpack, usemask=usemask, loose=loose, + invalid_raise=invalid_raise, max_rows=max_rows, encoding=encoding, + like=like + ) + if max_rows is not None: if skip_footer: raise ValueError( @@ -2250,6 +2304,11 @@ def genfromtxt(fname, dtype=float, comments='#', delimiter=None, return output.squeeze() +_genfromtxt_with_like = array_function_dispatch( + _genfromtxt_dispatcher +)(genfromtxt) + + def ndfromtxt(fname, **kwargs): """ Load ASCII data stored in a file and return it as a single array. diff --git a/numpy/lib/twodim_base.py b/numpy/lib/twodim_base.py index cd7484241..2b4cbdfbb 100644 --- a/numpy/lib/twodim_base.py +++ b/numpy/lib/twodim_base.py @@ -8,7 +8,7 @@ from numpy.core.numeric import ( asarray, where, int8, int16, int32, int64, empty, promote_types, diagonal, nonzero ) -from numpy.core.overrides import set_module +from numpy.core.overrides import set_array_function_like_doc, set_module from numpy.core import overrides from numpy.core import iinfo @@ -149,8 +149,13 @@ def flipud(m): return m[::-1, ...] +def _eye_dispatcher(N, M=None, k=None, dtype=None, order=None, *, like=None): + return (like,) + + +@set_array_function_like_doc @set_module('numpy') -def eye(N, M=None, k=0, dtype=float, order='C'): +def eye(N, M=None, k=0, dtype=float, order='C', *, like=None): """ Return a 2-D array with ones on the diagonal and zeros elsewhere. @@ -171,6 +176,9 @@ def eye(N, M=None, k=0, dtype=float, order='C'): column-major (Fortran-style) order in memory. .. versionadded:: 1.14.0 + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -194,6 +202,8 @@ def eye(N, M=None, k=0, dtype=float, order='C'): [0., 0., 0.]]) """ + if like is not None: + return _eye_with_like(N, M=M, k=k, dtype=dtype, order=order, like=like) if M is None: M = N m = zeros((N, M), dtype=dtype, order=order) @@ -207,6 +217,11 @@ def eye(N, M=None, k=0, dtype=float, order='C'): return m +_eye_with_like = array_function_dispatch( + _eye_dispatcher +)(eye) + + def _diag_dispatcher(v, k=None): return (v,) @@ -343,8 +358,13 @@ def diagflat(v, k=0): return wrap(res) +def _tri_dispatcher(N, M=None, k=None, dtype=None, *, like=None): + return (like,) + + +@set_array_function_like_doc @set_module('numpy') -def tri(N, M=None, k=0, dtype=float): +def tri(N, M=None, k=0, dtype=float, *, like=None): """ An array with ones at and below the given diagonal and zeros elsewhere. @@ -361,6 +381,9 @@ def tri(N, M=None, k=0, dtype=float): and `k` > 0 is above. The default is 0. dtype : dtype, optional Data type of the returned array. The default is float. + ${ARRAY_FUNCTION_LIKE} + + .. versionadded:: 1.20.0 Returns ------- @@ -381,6 +404,9 @@ def tri(N, M=None, k=0, dtype=float): [1., 1., 0., 0., 0.]]) """ + if like is not None: + return _tri_with_like(N, M=M, k=k, dtype=dtype, like=like) + if M is None: M = N @@ -393,6 +419,11 @@ def tri(N, M=None, k=0, dtype=float): return m +_tri_with_like = array_function_dispatch( + _tri_dispatcher +)(tri) + + def _trilu_dispatcher(m, k=None): return (m,) |