Repository: numpy/numpy
File: numpy/lib/tests/test_io.py (Python)

Code before:


class TestSaveLoad(RoundtripTest):
    def roundtrip(self, *args, **kwargs):
        RoundtripTest.roundtrip(self, np.save, *args, **kwargs)
        assert_equal(self.arr[0], self.arr_reloaded)
        assert_equal(self.arr[0].dtype, self.arr_reloaded.dtype)
        assert_equal(self.arr[0].flags.fnc, self.arr_reloaded.flags.fnc)


class TestSavezLoad(RoundtripTest):
    def roundtrip(self, *args, **kwargs):
        RoundtripTest.roundtrip(self, np.savez, *args, **kwargs)
        try:
            for n, arr in enumerate(self.arr):
                reloaded = self.arr_reloaded['arr_%d' % n]
                assert_equal(arr, reloaded)
                assert_equal(arr.dtype, reloaded.dtype)
                assert_equal(arr.flags.fnc, reloaded.flags.fnc)
        finally:
            # delete tempfile, must be done here on windows
            if self.arr_reloaded.fid:
                self.arr_reloaded.fid.close()
                os.remove(self.arr_reloaded.fid.name)

    def load_non_npy(self):
        """Test loading non-.npy files and name mapping in .npz."""
        with temppath(prefix="numpy_test_npz_load_non_npy_", suffix=".npz") as tmp:
            with zipfile.ZipFile(tmp) as npz:
                with npz.open("test1.npy", "w") as out_file:
                    np.save(out_file, np.arange(10))
                with npz.open("test2", "w") as out_file:
                    np.save(out_file, np.arange(10))
                with npz.open("metadata", "w") as out_file:
                    out_file.write(b"Name: Test")
            npz = np.load(tmp)
            assert len(npz["test1"]) == 10
            assert len(npz["test1.npy"]) == 10
            assert len(npz["test2"]) == 10
            assert npz["metadata"] == b"Name: Test"

    @pytest.mark.skipif(IS_PYPY, reason="Hangs on PyPy")
    @pytest.mark.skipif(not IS_64BIT, reason="Needs 64bit platform")
    @pytest.mark.slow
    def test_big_arrays(self):
        L = (1 << 31) + 100000
        a = np.empty(L, dtype=np.uint8)
        with temppath(prefix="numpy_test_big_arrays_", suffix=".npz") as tmp:
            np.savez(tmp, a=a)
            del a
            npfile = np.load(tmp)

Review comment (suggestion):
```suggestion
    def test_load_non_npy(self):
```

Code after:


class TestSaveLoad(RoundtripTest):
    def roundtrip(self, *args, **kwargs):
        RoundtripTest.roundtrip(self, np.save, *args, **kwargs)
        assert_equal(self.arr[0], self.arr_reloaded)
        assert_equal(self.arr[0].dtype, self.arr_reloaded.dtype)
        assert_equal(self.arr[0].flags.fnc, self.arr_reloaded.flags.fnc)


class TestSavezLoad(RoundtripTest):
    def roundtrip(self, *args, **kwargs):
        RoundtripTest.roundtrip(self, np.savez, *args, **kwargs)
        try:
            for n, arr in enumerate(self.arr):
                reloaded = self.arr_reloaded['arr_%d' % n]
                assert_equal(arr, reloaded)
                assert_equal(arr.dtype, reloaded.dtype)
                assert_equal(arr.flags.fnc, reloaded.flags.fnc)
        finally:
            # delete tempfile, must be done here on windows
            if self.arr_reloaded.fid:
                self.arr_reloaded.fid.close()
                os.remove(self.arr_reloaded.fid.name)

    def test_load_non_npy(self):
        """Test loading non-.npy files and name mapping in .npz."""
        with temppath(prefix="numpy_test_npz_load_non_npy_", suffix=".npz") as tmp:
            with zipfile.ZipFile(tmp, "w") as npz:
                with npz.open("test1.npy", "w") as out_file:
                    np.save(out_file, np.arange(10))
                with npz.open("test2", "w") as out_file:
                    np.save(out_file, np.arange(10))
                with npz.open("metadata", "w") as out_file:
                    out_file.write(b"Name: Test")
            with np.load(tmp) as npz:
                assert len(npz["test1"]) == 10
                assert len(npz["test1.npy"]) == 10
                assert len(npz["test2"]) == 10
                assert npz["metadata"] == b"Name: Test"

    @pytest.mark.skipif(IS_PYPY, reason="Hangs on PyPy")
    @pytest.mark.skipif(not IS_64BIT, reason="Needs 64bit platform")
    @pytest.mark.slow
    def test_big_arrays(self):
        L = (1 << 31) + 100000
        a = np.empty(L, dtype=np.uint8)
        with temppath(prefix="numpy_test_big_arrays_", suffix=".npz") as tmp:
            np.savez(tmp, a=a)
            del a
            npfile = np.load(tmp)