import pytest
import numpy as np
from numpy.testing import assert_array_equal
def test_matrix_transpose_raises_error_for_1d():
msg = "matrix transpose with ndim < 2 is undefined"
arr = np.arange(48)
with pytest.raises(ValueError, match=msg):
arr.mT
def test_matrix_transpose_equals_transpose_2d():
arr = np.arange(48).reshape((6, 8))
assert_array_equal(arr.T, arr.mT)
ARRAY_SHAPES_TO_TEST = (
(5, 2),
(5, 2, 3),
(5, 2, 3, 4),
)
@pytest.mark.parametrize("shape", ARRAY_SHAPES_TO_TEST)
def test_matrix_transpose_equals_swapaxes(shape):
num_of_axes = len(shape)
vec = np.arange(shape[-1])
arr = np.broadcast_to(vec, shape)
tgt = np.swapaxes(arr, num_of_axes - 2, num_of_axes - 1)
mT = arr.mT
assert_array_equal(tgt, mT)
class MyArr(np.ndarray):
def __array_wrap__(self, arr, context=None, return_scalar=None):
return super().__array_wrap__(arr, context, return_scalar)
class MyArrNoWrap(np.ndarray):
pass
@pytest.mark.parametrize("subclass_self", [np.ndarray, MyArr, MyArrNoWrap])
@pytest.mark.parametrize("subclass_arr", [np.ndarray, MyArr, MyArrNoWrap])
def test_array_wrap(subclass_self, subclass_arr):
# NumPy should allow `__array_wrap__` to be called on arrays, it's logic
# is designed in a way that:
#
# * Subclasses never return scalars by default (to preserve their
# information). They can choose to if they wish.
# * NumPy returns scalars, if `return_scalar` is passed as True to allow
# manual calls to `arr.__array_wrap__` to do the right thing.
# * The type of the input should be ignored (it should be a base-class
# array, but I am not sure this is guaranteed).
arr = np.arange(3).view(subclass_self)
arr0d = np.array(3, dtype=np.int8).view(subclass_arr)
# With third argument True, ndarray allows "decay" to scalar.
# (I don't think NumPy would pass `None`, but it seems clear to support)
if subclass_self is np.ndarray:
assert type(arr.__array_wrap__(arr0d, None, True)) is np.int8
else:
assert type(arr.__array_wrap__(arr0d, None, True)) is type(arr)
# Otherwise, result should be viewed as the subclass
assert type(arr.__array_wrap__(arr0d)) is type(arr)
assert type(arr.__array_wrap__(arr0d, None, None)) is type(arr)
assert type(arr.__array_wrap__(arr0d, None, False)) is type(arr)
# Non 0-D array can't be converted to scalar, so we ignore that
arr1d = np.array([3], dtype=np.int8).view(subclass_arr)
assert type(arr.__array_wrap__(arr1d, None, True)) is type(arr)