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

Code before:
        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]
        assert not non_contiguous_array.flags.c_contiguous
        assert not non_contiguous_array.flags.f_contiguous

        # make sure non-contiguous arrays can be pickled-depickled
        # using any protocol
        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
            depickled_non_contiguous_array = pickle.loads(
                    pickle.dumps(non_contiguous_array, protocol=proto))

            assert_equal(non_contiguous_array, depickled_non_contiguous_array)

    def test_roundtrip(self):
        for proto in range(2, pickle.HIGHEST_PROTOCOL + 1):
            carray = np.array([[2, 9], [7, 0], [3, 8]])
            DATA = [
                carray,
                np.transpose(carray),
                np.array([('xxx', 1, 2.0)], dtype=[('a', (str, 3)), ('b', int),
                                                   ('c', float)])

Review comment (suggestion):
Thanks, this works (I suspect you could use `dtype=np.int8` or so and just copy the whole bytes string into the code here also), but just if you like it.

Code after:

    @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

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

        assert_equal(transposed_contiguous_array, depickled_transposed_contiguous_array)

    @pytest.mark.skipif(pickle.HIGHEST_PROTOCOL < 5, reason="requires pickle protocol 5")
    def test_load_legacy_pkl_protocol5(self):
        # legacy byte strs are dumped in 2.2.1
        c_contiguous_dumped = b'\x80\x05\x95\x90\x00\x00\x00\x00\x00\x00\x00\x8c\x13numpy._core.numeric\x94\x8c\x0b_frombuffer\x94\x93\x94(\x96\x18\x00\x00\x00\x00\x00\x00\x00\x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x94\x8c\x05numpy\x94\x8c\x05dtype\x94\x93\x94\x8c\x02u1\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94bK\x03K\x04K\x02\x87\x94\x8c\x01C\x94t\x94R\x94.' # noqa
        f_contiguous_dumped = b'\x80\x05\x95\x90\x00\x00\x00\x00\x00\x00\x00\x8c\x13numpy._core.numeric\x94\x8c\x0b_frombuffer\x94\x93\x94(\x96\x18\x00\x00\x00\x00\x00\x00\x00\x00\x01\x02\x03\x04\x05\x06\x07\x08\t\n\x0b\x0c\r\x0e\x0f\x10\x11\x12\x13\x14\x15\x16\x17\x94\x8c\x05numpy\x94\x8c\x05dtype\x94\x93\x94\x8c\x02u1\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94bK\x03K\x04K\x02\x87\x94\x8c\x01F\x94t\x94R\x94.' # noqa
        transposed_contiguous_dumped = b'\x80\x05\x95\xa5\x00\x00\x00\x00\x00\x00\x00\x8c\x16numpy._core.multiarray\x94\x8c\x0c_reconstruct\x94\x93\x94\x8c\x05numpy\x94\x8c\x07ndarray\x94\x93\x94K\x00\x85\x94C\x01b\x94\x87\x94R\x94(K\x01K\x04K\x03K\x02\x87\x94h\x03\x8c\x05dtype\x94\x93\x94\x8c\x02u1\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94b\x89C\x18\x00\x01\x08\t\x10\x11\x02\x03\n\x0b\x12\x13\x04\x05\x0c\r\x14\x15\x06\x07\x0e\x0f\x16\x17\x94t\x94b.' # noqa
        no_contiguous_dumped = b'\x80\x05\x95\x91\x00\x00\x00\x00\x00\x00\x00\x8c\x16numpy._core.multiarray\x94\x8c\x0c_reconstruct\x94\x93\x94\x8c\x05numpy\x94\x8c\x07ndarray\x94\x93\x94K\x00\x85\x94C\x01b\x94\x87\x94R\x94(K\x01K\x03K\x02\x86\x94h\x03\x8c\x05dtype\x94\x93\x94\x8c\x02u1\x94\x89\x88\x87\x94R\x94(K\x03\x8c\x01|\x94NNNJ\xff\xff\xff\xffJ\xff\xff\xff\xffK\x00t\x94b\x89C\x06\x00\x01\x04\x05\x08\t\x94t\x94b.' # noqa
        x = np.arange(24, dtype='uint8').reshape(3, 4, 2)
        assert_equal(x, pickle.loads(c_contiguous_dumped))
        x = np.arange(24, dtype='uint8').reshape(3, 4, 2, order='F')
        assert_equal(x, pickle.loads(f_contiguous_dumped))
        x = np.arange(24, dtype='uint8').reshape(3, 4, 2).transpose((1, 0, 2))
        assert_equal(x, pickle.loads(transposed_contiguous_dumped))