Clean up MemoryFile /vsimem/ dirs recursively

CHANGES.txt

@@ -6,6 +6,7 @@ Changes
 
 Bug fixes:
 
+- Better, recursive cleanup of MemoryFile /vsimem/ directories.
 - Python ignores SIGPIPE by default. By never catching BrokenPipeError via
   `except Exception` when, for example, piping the output of rio-shapes to
   the Unix head program, we avoid getting an unhandled BrokenPipeError message

rasterio/_env.pyx

@@ -450,6 +450,14 @@ cdef class GDALEnv(ConfigEnv):
     def _dump_open_datasets(self):
         GDALDumpOpenDatasets(stderr)
 
+    def _dump_vsimem(self):
+        dirs = VSIReadDir("/vsimem/")
+        num_dirs = CSLCount(dirs)
+        try:
+            return list([dirs[i] for i in range(num_dirs) if str(dirs[i])])
+        finally:
+            CSLDestroy(dirs)
+
 
 def set_proj_data_search_path(path):
     """Set PROJ data search path"""

rasterio/_io.pyx

@@ -1226,8 +1226,7 @@ cdef class MemoryFileBase:
         if self._vsif != NULL:
             VSIFCloseL(self._vsif)
         self._vsif = NULL
-        _delete_dataset_if_exists(self.name)
-        VSIRmdir(self._dirname.encode("utf-8"))
+        VSIRmdirRecursive("/vsimem/{}".format(self._dirname).encode("utf-8"))
         self.closed = True
 
     def seek(self, offset, whence=0):

rasterio/env.py

@@ -261,6 +261,13 @@ class Env:
         """
         return local._env._dump_open_datasets()
 
+    def _dump_vsimem(self):
+        """Returns contents of /vsimem/.
+
+        For debugging and testing purposes.
+        """
+        return local._env._dump_vsimem()
+
     def __enter__(self):
         log.debug("Entering env context: %r", self)
         if local._env is None:

rasterio/gdal.pxi

@@ -131,7 +131,9 @@ cdef extern from "cpl_vsi.h" nogil:
     int VSIFCloseL(VSILFILE *fp)
     int VSIUnlink(const char *path)
     int VSIMkdir(const char *path, long mode)
+    char** VSIReadDir(const char *path)
     int VSIRmdir(const char *path)
+    int VSIRmdirRecursive(const char *path)
     int VSIFFlushL(VSILFILE *fp)
     size_t VSIFReadL(void *buffer, size_t nSize, size_t nCount, VSILFILE *fp)
     int VSIFSeekL(VSILFILE *fp, vsi_l_offset nOffset, int nWhence)

tests/test_memoryfile.py

@@ -4,6 +4,7 @@ Tests in this file will ONLY run for GDAL >= 2.x"""
 from io import BytesIO
 import logging
 import os.path
+from pathlib import Path
 
 from affine import Affine
 import numpy
@@ -299,15 +300,29 @@ def test_write_plus_mode():
 
 
 def test_write_plus_model_jpeg():
-    with rasterio.Env(), MemoryFile() as memfile:
-        with memfile.open(driver='JPEG', dtype='uint8', count=3, height=32, width=32, crs='epsg:3226', transform=Affine.identity() * Affine.scale(0.5, -0.5)) as dst:
-            dst.write(numpy.full((32, 32), 255, dtype='uint8'), 1)
-            dst.write(numpy.full((32, 32), 204, dtype='uint8'), 2)
-            dst.write(numpy.full((32, 32), 153, dtype='uint8'), 3)
-            data = dst.read()
-            assert (data[0] == 255).all()
-            assert (data[1] == 204).all()
-            assert (data[2] == 153).all()
+    """Ensure /vsimem/ file is cleaned up."""
+    with rasterio.Env() as env:
+        with MemoryFile() as memfile:
+            with memfile.open(
+                driver="JPEG",
+                dtype="uint8",
+                count=3,
+                height=32,
+                width=32,
+                crs="epsg:3226",
+                transform=Affine.identity() * Affine.scale(0.5, -0.5),
+            ) as dst:
+                dst.write(numpy.full((32, 32), 255, dtype="uint8"), 1)
+                dst.write(numpy.full((32, 32), 204, dtype="uint8"), 2)
+                dst.write(numpy.full((32, 32), 153, dtype="uint8"), 3)
+                data = dst.read()
+                assert (data[0] == 255).all()
+                assert (data[1] == 204).all()
+                assert (data[2] == 153).all()
+
+        assert (
+            Path(*Path(memfile.name).parts[2:]).parent.as_posix() not in env._dump_vsimem()
+        )
 
 
 def test_memfile_copyfiles(path_rgb_msk_byte_tif):