API: Introduce numpy.astype

mtsokol · mtsokol · commit adf2d3b4000d · 2023-11-06T15:55:50.000+01:00
diff --git a/doc/source/reference/routines.array-creation.rst b/doc/source/reference/routines.array-creation.rst
@@ -33,6 +33,7 @@ From existing data
    asanyarray
    ascontiguousarray
    asmatrix
+   astype
    copy
    frombuffer
    from_dlpack
diff --git a/numpy/__init__.py b/numpy/__init__.py
@@ -126,7 +126,7 @@
         arctan2, arctanh, argmax, argmin, argpartition, argsort, argwhere,
         around, array, array2string, array_equal, array_equiv, array_repr,
         array_str, asanyarray, asarray, ascontiguousarray, asfortranarray,
-        atleast_1d, atleast_2d, atleast_3d, base_repr, binary_repr,
+        astype, atleast_1d, atleast_2d, atleast_3d, base_repr, binary_repr,
         bitwise_and, bitwise_count, bitwise_not, bitwise_or, bitwise_xor,
         block, bool_, broadcast, busday_count, busday_offset, busdaycalendar,
         byte, bytes_, can_cast, cbrt, cdouble, ceil, character, choose, clip,
diff --git a/numpy/__init__.pyi b/numpy/__init__.pyi
@@ -376,6 +376,7 @@ from numpy._core.numeric import (
     isclose as isclose,
     array_equal as array_equal,
     array_equiv as array_equiv,
+    astype as astype,
 )
 
 from numpy._core.numerictypes import (
diff --git a/numpy/_core/numeric.py b/numpy/_core/numeric.py
@@ -9,16 +9,15 @@
 import numpy as np
 from . import multiarray
 from .multiarray import (
-    ALLOW_THREADS,
-    BUFSIZE, CLIP, MAXDIMS, MAY_SHARE_BOUNDS, MAY_SHARE_EXACT, RAISE,
-    WRAP, arange, array, asarray, asanyarray, ascontiguousarray,
-    asfortranarray, broadcast, can_cast,
-    concatenate, copyto, dot, dtype, empty,
-    empty_like, flatiter, frombuffer, from_dlpack, fromfile, fromiter,
-    fromstring, inner, lexsort, matmul, may_share_memory,
-    min_scalar_type, ndarray, nditer, nested_iters, promote_types,
-    putmask, result_type, shares_memory, vdot, where,
-    zeros, normalize_axis_index, _get_promotion_state, _set_promotion_state)
+    ALLOW_THREADS, BUFSIZE, CLIP, MAXDIMS, MAY_SHARE_BOUNDS, MAY_SHARE_EXACT,
+    RAISE, WRAP, arange, array, asarray, asanyarray, ascontiguousarray,
+    asfortranarray, broadcast, can_cast, concatenate, copyto, dot, dtype,
+    empty, empty_like, flatiter, frombuffer, from_dlpack, fromfile, fromiter,
+    fromstring, inner, lexsort, matmul, may_share_memory, min_scalar_type,
+    ndarray, nditer, nested_iters, promote_types, putmask, result_type,
+    shares_memory, vdot, where, zeros, normalize_axis_index,
+    _get_promotion_state, _set_promotion_state
+)
 
 from . import overrides
 from . import umath
@@ -43,7 +42,7 @@
     'arange', 'array', 'asarray', 'asanyarray', 'ascontiguousarray',
     'asfortranarray', 'zeros', 'count_nonzero', 'empty', 'broadcast', 'dtype',
     'fromstring', 'fromfile', 'frombuffer', 'from_dlpack', 'where',
-    'argwhere', 'copyto', 'concatenate', 'lexsort',
+    'argwhere', 'copyto', 'concatenate', 'lexsort', 'astype',
     'can_cast', 'promote_types', 'min_scalar_type',
     'result_type', 'isfortran', 'empty_like', 'zeros_like', 'ones_like',
     'correlate', 'convolve', 'inner', 'dot', 'outer', 'vdot', 'roll',
@@ -2508,6 +2507,60 @@ def array_equiv(a1, a2):
     return bool(asarray(a1 == a2).all())
 
 
+def _astype_dispatcher(x, dtype, /, *, copy=None):
+    return (x, dtype)
+
+
+@array_function_dispatch(_astype_dispatcher)
+def astype(x, dtype, /, *, copy = True):
+    """
+    Copies an array to a specified data type.
+
+    This function is an Array API compatible alternative to
+    `numpy.ndarray.astype`.
+
+    Parameters
+    ----------
+    x : array_like
+        Input array to cast.
+    dtype : dtype
+        Data type of the result.
+    copy : bool, optional
+        Specifies whether to copy an array when the specified dtype matches
+        the data type of the input array ``x``. If ``True``, a newly allocated
+        array must always be returned. If ``False`` and the specified dtype
+        matches the data type of the input array, the input array must be
+        returned; otherwise, a newly allocated array must be returned.
+        Defaults to ``True``.
+
+    Returns
+    -------
+    out : ndarray
+        An array having the specified data type.
+
+    See Also
+    --------
+    ndarray.astype
+
+    Examples
+    --------
+    >>> arr = np.array([1, 2, 3]); arr
+    array([1, 2, 3])
+    >>> np.astype(arr, np.float64)
+    array([1., 2., 3.])
+
+    Non-copy case:
+
+    >>> arr = np.array([1, 2, 3])
+    >>> arr_noncpy = np.astype(arr, arr.dtype, copy=False)
+    >>> np.shares_memory(arr, arr_noncpy)
+    True
+
+    """
+    x = asarray(x)
+    return x.astype(dtype, copy=copy)
+
+
 inf = PINF
 nan = NAN
 False_ = bool_(False)
diff --git a/numpy/_core/numeric.pyi b/numpy/_core/numeric.pyi
@@ -658,3 +658,9 @@ def isclose(
 def array_equal(a1: ArrayLike, a2: ArrayLike, equal_nan: bool = ...) -> bool: ...
 
 def array_equiv(a1: ArrayLike, a2: ArrayLike) -> bool: ...
+
+def astype(
+    x: ArrayLike,
+    dtype: DTypeLike,
+    copy: bool = ...,
+) -> NDArray[Any]: ...
diff --git a/numpy/_core/tests/test_numeric.py b/numpy/_core/tests/test_numeric.py
@@ -3981,3 +3981,20 @@ def test_zero_dimensional(self):
         arr_0d = np.array(1)
         ret = np.tensordot(arr_0d, arr_0d, ([], []))  # contracting no axes is well defined
         assert_array_equal(ret, arr_0d)
+
+
+class TestAsType:
+
+    def test_astype(self):
+        data = [[1, 2], [3, 4]]
+        arr = np.astype(
+            np.array(data, dtype=np.int64), np.uint32
+        )
+        expected = np.array(data, dtype=np.uint32)
+
+        assert_array_equal(arr, expected)
+        assert_equal(arr.dtype, expected.dtype)
+
+        assert np.shares_memory(
+            arr, np.astype(arr, arr.dtype, copy=False)
+        )

Original file line number	Diff line number	Diff line change
`@@ -376,6 +376,7 @@ from numpy._core.numeric import (`
`376`	`376`	`isclose as isclose,`
`377`	`377`	`array_equal as array_equal,`
`378`	`378`	`array_equiv as array_equiv,`
	`379`	`+ astype as astype,`
`379`	`380`	`)`
`380`	`381`
`381`	`382`	`from numpy._core.numerictypes import (`