diff --git a/CHANGES.txt b/CHANGES.txt index e590a6aa..3337bf94 100644 --- a/CHANGES.txt +++ b/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 diff --git a/rasterio/_env.pyx b/rasterio/_env.pyx index e970685f..9988aa1a 100644 --- a/rasterio/_env.pyx +++ b/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""" diff --git a/rasterio/_io.pyx b/rasterio/_io.pyx index f84c9a35..ed45dddd 100644 --- a/rasterio/_io.pyx +++ b/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): diff --git a/rasterio/env.py b/rasterio/env.py index 5de53bea..ae0c54de 100644 --- a/rasterio/env.py +++ b/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: diff --git a/rasterio/gdal.pxi b/rasterio/gdal.pxi index a9dcaf55..113e9dbe 100644 --- a/rasterio/gdal.pxi +++ b/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) diff --git a/tests/test_memoryfile.py b/tests/test_memoryfile.py index 43ab03c5..c7536c62 100644 --- a/tests/test_memoryfile.py +++ b/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):