codekingpro/portable-devtools

1import sys
2
3import pytest
4
5import numpy as np
6from numpy.testing import HAS_REFCOUNT, assert_array_equal
7
8
9def test_matrix_transpose_raises_error_for_1d():
10    msg = "matrix transpose with ndim < 2 is undefined"
11    arr = np.arange(48)
12    with pytest.raises(ValueError, match=msg):
13        arr.mT
14
15
16def test_matrix_transpose_equals_transpose_2d():
17    arr = np.arange(48).reshape((6, 8))
18    assert_array_equal(arr.T, arr.mT)
19
20
21ARRAY_SHAPES_TO_TEST = (
22    (5, 2),
23    (5, 2, 3),
24    (5, 2, 3, 4),
25)
26
27
28@pytest.mark.parametrize("shape", ARRAY_SHAPES_TO_TEST)
29def test_matrix_transpose_equals_swapaxes(shape):
30    num_of_axes = len(shape)
31    vec = np.arange(shape[-1])
32    arr = np.broadcast_to(vec, shape)
33    tgt = np.swapaxes(arr, num_of_axes - 2, num_of_axes - 1)
34    mT = arr.mT
35    assert_array_equal(tgt, mT)
36
37
38class MyArr(np.ndarray):
39    def __array_wrap__(self, arr, context=None, return_scalar=None):
40        return super().__array_wrap__(arr, context, return_scalar)
41
42
43class MyArrNoWrap(np.ndarray):
44    pass
45
46
47@pytest.mark.parametrize("subclass_self", [np.ndarray, MyArr, MyArrNoWrap])
48@pytest.mark.parametrize("subclass_arr", [np.ndarray, MyArr, MyArrNoWrap])
49def test_array_wrap(subclass_self, subclass_arr):
50    # NumPy should allow `__array_wrap__` to be called on arrays, it's logic
51    # is designed in a way that:
52    #
53    # * Subclasses never return scalars by default (to preserve their
54    #   information).  They can choose to if they wish.
55    # * NumPy returns scalars, if `return_scalar` is passed as True to allow
56    #   manual calls to `arr.__array_wrap__` to do the right thing.
57    # * The type of the input should be ignored (it should be a base-class
58    #   array, but I am not sure this is guaranteed).
59
60    arr = np.arange(3).view(subclass_self)
61
62    arr0d = np.array(3, dtype=np.int8).view(subclass_arr)
63    # With third argument True, ndarray allows "decay" to scalar.
64    # (I don't think NumPy would pass `None`, but it seems clear to support)
65    if subclass_self is np.ndarray:
66        assert type(arr.__array_wrap__(arr0d, None, True)) is np.int8
67    else:
68        assert type(arr.__array_wrap__(arr0d, None, True)) is type(arr)
69
70    # Otherwise, result should be viewed as the subclass
71    assert type(arr.__array_wrap__(arr0d)) is type(arr)
72    assert type(arr.__array_wrap__(arr0d, None, None)) is type(arr)
73    assert type(arr.__array_wrap__(arr0d, None, False)) is type(arr)
74
75    # Non 0-D array can't be converted to scalar, so we ignore that
76    arr1d = np.array([3], dtype=np.int8).view(subclass_arr)
77    assert type(arr.__array_wrap__(arr1d, None, True)) is type(arr)
78
79
80@pytest.mark.skipif(not HAS_REFCOUNT, reason="Python lacks refcounts")
81def test_cleanup_with_refs_non_contig():
82    # Regression test, leaked the dtype (but also good for rest)
83    dtype = np.dtype("O,i")
84    obj = object()
85    expected_ref_dtype = sys.getrefcount(dtype)
86    expected_ref_obj = sys.getrefcount(obj)
87    proto = np.full((3, 4, 5, 6, 7), np.array((obj, 2), dtype=dtype))
88    # Give array a non-trivial order to exercise more cleanup paths.
89    arr = proto.transpose((2, 0, 3, 1, 4)).copy("K")
90    del proto, arr
91
92    actual_ref_dtype = sys.getrefcount(dtype)
93    actual_ref_obj = sys.getrefcount(obj)
94    assert actual_ref_dtype == expected_ref_dtype
95    assert actual_ref_obj == actual_ref_dtype
96