Repository: numpy/numpy
File: numpy/_core/tests/test_multiarray.py (Python)

Code before:

            assert_equal(arr_with_object.dtype,
                         depickled_arr_with_object.dtype)
            assert_equal(arr_without_object.dtype,
                         depickled_arr_without_object.dtype)

    @pytest.mark.skipif(pickle.HIGHEST_PROTOCOL < 5,
                        reason="requires pickle protocol 5")
    def test_f_contiguous_array(self):
        f_contiguous_array = np.array([[1, 2, 3], [4, 5, 6]], order='F')
        buffers = []

        # When using pickle protocol 5, Fortran-contiguous arrays can be
        # serialized using out-of-band buffers
        bytes_string = pickle.dumps(f_contiguous_array, protocol=5,
                                    buffer_callback=buffers.append)

        assert len(buffers) > 0

        depickled_f_contiguous_array = pickle.loads(bytes_string,
                                                    buffers=buffers)

        assert_equal(f_contiguous_array, depickled_f_contiguous_array)
    
    def test_transposed_contiguous_array(self):
        transposed_contiguous_array = np.random.rand(2, 3, 4).transpose((1, 0, 2))
        buffers = []

        # When using pickle protocol 5, Fortran-contiguous arrays can be
        # serialized using out-of-band buffers
        bytes_string = pickle.dumps(transposed_contiguous_array, protocol=5,
                                    buffer_callback=buffers.append)

        assert len(buffers) > 0

        depickled_transposed_contiguous_array = pickle.loads(bytes_string,
                                                    buffers=buffers)

        assert_equal(transposed_contiguous_array, depickled_transposed_contiguous_array)
        
    def test_load_legacy_pkl(self):
        x = np.arange(24).reshape(3, 4, 2)
        # generated by same code in 2.2.1
        data_dir = os.path.join(os.path.dirname(__file__), 'data')
        legacy_pkl_path = os.path.join(data_dir, "legacy_contiguous.pkl")
        with open(legacy_pkl_path, 'rb') as f:
            y = pickle.load(f)
        assert_allclose(x, y)

    def test_non_contiguous_array(self):
        non_contiguous_array = np.arange(12).reshape(3, 4)[:, :2]

Review comment (suggestion):
The transpose example is not chosen carefully, maybe you should just have multiple examples with `pytest.parametrize` and make at least one 4-5 dimensional so this is less likely to happen:
```
a = np.ones((1, 2, 3))
a.transpose(1, 0, 2).transpose(1, 0, 2).shape == a.shape
```
so I don't know from the test if the back-transpose is filled correctly.

Code after:
        my_object = object()

        arr_with_object = np.array(
                [(my_object, 1, 2.0)],
                dtype=[('a', object), ('b', int), ('c', float)])
        arr_without_object = np.array(
                [('xxx', 1, 2.0)],
                dtype=[('a', str), ('b', int), ('c', float)])

        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
            depickled_arr_with_object = pickle.loads(
                    pickle.dumps(arr_with_object, protocol=proto))
            depickled_arr_without_object = pickle.loads(
                    pickle.dumps(arr_without_object, protocol=proto))

            assert_equal(arr_with_object.dtype,
                         depickled_arr_with_object.dtype)
            assert_equal(arr_without_object.dtype,
                         depickled_arr_without_object.dtype)

    @pytest.mark.skipif(pickle.HIGHEST_PROTOCOL < 5,
                        reason="requires pickle protocol 5")
    def test_f_contiguous_array(self):
        f_contiguous_array = np.array([[1, 2, 3], [4, 5, 6]], order='F')
        buffers = []

        # When using pickle protocol 5, Fortran-contiguous arrays can be
        # serialized using out-of-band buffers
        bytes_string = pickle.dumps(f_contiguous_array, protocol=5,
                                    buffer_callback=buffers.append)

        assert len(buffers) > 0

        depickled_f_contiguous_array = pickle.loads(bytes_string,
                                                    buffers=buffers)

        assert_equal(f_contiguous_array, depickled_f_contiguous_array)

    @pytest.mark.skipif(pickle.HIGHEST_PROTOCOL < 5, reason="requires pickle protocol 5")
    @pytest.mark.parametrize('transposed_contiguous_array', 
        [np.random.default_rng(42).random((2, 3, 4)).transpose((1, 0, 2)), 
         np.random.default_rng(42).random((2, 3, 4, 5)).transpose((1, 3, 0, 2))] + 
        [np.random.default_rng(42).random(np.arange(2, 7)).transpose(np.random.permutation(5)) for _ in range(3)])
    def test_transposed_contiguous_array(self, transposed_contiguous_array):
        buffers = []
        # When using pickle protocol 5, arrays which can be transposed to c_contiguous
        # can be serialized using out-of-band buffers
        bytes_string = pickle.dumps(transposed_contiguous_array, protocol=5,
                                    buffer_callback=buffers.append)

        assert len(buffers) > 0