Repository: scipy/scipy
File: scipy/sparse/tests/test_matrix_io.py (Python)

Code before:

    #save/load array
    fd, tmpfile = tempfile.mkstemp(suffix='.npz')
    os.close(fd)
    try:
        save_npz(tmpfile, csr_array([[1.2, 0, 0.9], [0, 0.3, 0]]))
        loaded_array = load_npz(tmpfile)
    finally:
        os.remove(tmpfile)

    assert not isinstance(loaded_matrix, sparray)
    assert isinstance(loaded_array, sparray)
    assert_(loaded_matrix.dtype == loaded_array.dtype)
    assert_equal(loaded_matrix.toarray(), loaded_array.toarray())

def test_nd_coo_format():
    A = coo_array([[[0]]])

    #save/load array
    fd, tmpfile = tempfile.mkstemp(suffix='.npz')
    os.close(fd)
    try:
        save_npz(tmpfile, A)
        loaded_A = load_npz(tmpfile)
    finally:
        os.remove(tmpfile)

    assert isinstance(loaded_A, coo_array)
    assert_(loaded_A.shape == A.shape)
    assert_equal(A.toarray(), loaded_A.toarray())

def test_malicious_load():
    class Executor:
        def __reduce__(self):
            return (assert_, (False, 'unexpected code execution'))

    fd, tmpfile = tempfile.mkstemp(suffix='.npz')
    os.close(fd)
    try:
        np.savez(tmpfile, format=Executor())

        # Should raise a ValueError, not execute code
        assert_raises(ValueError, load_npz, tmpfile)
    finally:
        os.remove(tmpfile)


def test_py23_compatibility():
    # Try loading files saved on Python 2 and Python 3.  They are not
    # the same, since files saved with SciPy versions < 1.0.0 may
    # contain unicode.

Review comment (bug):
Maybe we could cut down on the boilerplate for file handling here -- the diff below seems to work locally. I did also confirm that the test fails fails before/passes after the patch while I was here.

```diff
-def test_nd_coo_format():
+def test_nd_coo_format(tmpdir):
     A = coo_array([[[0]]])
 
     #save/load array
-    fd, tmpfile = tempfile.mkstemp(suffix='.npz')
-    os.close(fd)
-    try:
+    with tmpdir.as_cwd():
+        tmpfile = "f.npz"
         save_npz(tmpfile, A)
         loaded_A = load_npz(tmpfile)
-    finally:
-        os.remove(tmpfile)
```

Code after:
        os.remove(tmpfile)

    #save/load array
    fd, tmpfile = tempfile.mkstemp(suffix='.npz')
    os.close(fd)
    try:
        save_npz(tmpfile, csr_array([[1.2, 0, 0.9], [0, 0.3, 0]]))
        loaded_array = load_npz(tmpfile)
    finally:
        os.remove(tmpfile)

    assert not isinstance(loaded_matrix, sparray)
    assert isinstance(loaded_array, sparray)
    assert_(loaded_matrix.dtype == loaded_array.dtype)
    assert_equal(loaded_matrix.toarray(), loaded_array.toarray())

@pytest.mark.parametrize("value", [0, 1.2])
@pytest.mark.parametrize("ndim", [1, 2, 3])
def test_nd_coo_format(ndim, value, tmpdir):
    A = coo_array([value]).reshape((1,) * ndim)

    #save/load array
    with tmpdir.as_cwd():
        tmpfile = "f.npz"

        save_npz(tmpfile, A)
        loaded_A = load_npz(tmpfile)

    assert isinstance(loaded_A, coo_array)
    assert_(loaded_A.shape == A.shape)
    assert_equal(A.toarray(), loaded_A.toarray())

def test_malicious_load():
    class Executor:
        def __reduce__(self):
            return (assert_, (False, 'unexpected code execution'))

    fd, tmpfile = tempfile.mkstemp(suffix='.npz')
    os.close(fd)
    try:
        np.savez(tmpfile, format=Executor())

        # Should raise a ValueError, not execute code
        assert_raises(ValueError, load_npz, tmpfile)
    finally:
        os.remove(tmpfile)


def test_py23_compatibility():
    # Try loading files saved on Python 2 and Python 3.  They are not
    # the same, since files saved with SciPy versions < 1.0.0 may