1"""MemoryFile tests. MemoryFile requires GDAL 2.0+. 2Tests in this file will ONLY run for GDAL >= 2.x""" 3 4from io import BytesIO 5import logging 6import os.path 7 8from affine import Affine 9import numpy 10import pytest 11 12import rasterio 13from rasterio.io import MemoryFile, ZipMemoryFile 14from rasterio.enums import MaskFlags 15from rasterio.env import GDALVersion 16from rasterio.shutil import copyfiles 17 18 19# Skip ENTIRE module if not GDAL >= 2.x. 20# pytestmark is a keyword that instructs pytest to skip this module. 21pytestmark = pytest.mark.skipif( 22 not GDALVersion.runtime().major >= 2, 23 reason="MemoryFile requires GDAL 2.x") 24 25 26@pytest.fixture(scope='session') 27def rgb_file_bytes(path_rgb_byte_tif): 28 """Get the bytes of our RGB.bytes.tif file""" 29 return open(path_rgb_byte_tif, 'rb').read() 30 31 32@pytest.fixture(scope='session') 33def rgb_lzw_file_bytes(): 34 """Get the bytes of our RGB.bytes.tif file""" 35 return open('tests/data/rgb_lzw.tif', 'rb').read() 36 37 38@pytest.fixture(scope='function') 39def rgb_file_object(path_rgb_byte_tif): 40 """Get RGB.bytes.tif file opened in 'rb' mode""" 41 return open(path_rgb_byte_tif, 'rb') 42 43 44@pytest.fixture(scope='session') 45def rgb_data_and_profile(path_rgb_byte_tif): 46 with rasterio.open(path_rgb_byte_tif) as src: 47 data = src.read() 48 profile = src.profile 49 return data, profile 50 51 52def test_initial_empty(): 53 with MemoryFile() as memfile: 54 assert len(memfile) == 0 55 assert len(memfile.getbuffer()) == 0 56 assert memfile.tell() == 0 57 58 59def test_initial_not_bytes(): 60 """Creating a MemoryFile from not bytes fails.""" 61 with pytest.raises(TypeError): 62 MemoryFile(u'lolwut') 63 64 65def test_initial_bytes(rgb_file_bytes): 66 """MemoryFile contents can initialized from bytes and opened.""" 67 with MemoryFile(rgb_file_bytes) as memfile: 68 with memfile.open() as src: 69 assert src.driver == 'GTiff' 70 assert src.count == 3 71 assert src.dtypes == ('uint8', 'uint8', 'uint8') 72 assert src.read().shape == (3, 718, 791) 73 74 75def test_initial_lzw_bytes(rgb_lzw_file_bytes): 76 """MemoryFile contents can initialized from bytes and opened.""" 77 with MemoryFile(rgb_lzw_file_bytes) as memfile: 78 with memfile.open() as src: 79 assert src.driver == 'GTiff' 80 assert src.count == 3 81 assert src.dtypes == ('uint8', 'uint8', 'uint8') 82 assert src.read().shape == (3, 718, 791) 83 84 85def test_initial_file_object(rgb_file_object): 86 """MemoryFile contents can initialized from bytes and opened.""" 87 with MemoryFile(rgb_file_object) as memfile: 88 with memfile.open() as src: 89 assert src.driver == 'GTiff' 90 assert src.count == 3 91 assert src.dtypes == ('uint8', 'uint8', 'uint8') 92 assert src.read().shape == (3, 718, 791) 93 94 95def test_closed(): 96 """A closed MemoryFile can not be opened""" 97 with MemoryFile() as memfile: 98 pass 99 with pytest.raises(IOError): 100 memfile.open() 101 102 103def test_non_initial_bytes(rgb_file_bytes): 104 """MemoryFile contents can be read from bytes and opened.""" 105 with MemoryFile() as memfile: 106 assert memfile.write(rgb_file_bytes) == len(rgb_file_bytes) 107 with memfile.open() as src: 108 assert src.driver == 'GTiff' 109 assert src.count == 3 110 assert src.dtypes == ('uint8', 'uint8', 'uint8') 111 assert src.read().shape == (3, 718, 791) 112 113 114def test_non_initial_bytes_in_two(rgb_file_bytes): 115 """MemoryFile contents can be read from bytes in two steps and opened.""" 116 with MemoryFile() as memfile: 117 assert memfile.write(rgb_file_bytes[:10]) == 10 118 assert memfile.write(rgb_file_bytes[10:]) == len(rgb_file_bytes) - 10 119 with memfile.open() as src: 120 assert src.driver == 'GTiff' 121 assert src.count == 3 122 assert src.dtypes == ('uint8', 'uint8', 'uint8') 123 assert src.read().shape == (3, 718, 791) 124 125 126def test_non_initial_bytes_in_two_reverse(rgb_file_bytes): 127 """MemoryFile contents can be read from bytes in two steps, tail first, and opened. 128 Demonstrates fix of #1926.""" 129 with MemoryFile() as memfile: 130 memfile.seek(600000) 131 assert memfile.write(rgb_file_bytes[600000:]) == len(rgb_file_bytes) - 600000 132 memfile.seek(0) 133 assert memfile.write(rgb_file_bytes[:600000]) == 600000 134 with memfile.open() as src: 135 assert src.driver == "GTiff" 136 assert src.count == 3 137 assert src.dtypes == ("uint8", "uint8", "uint8") 138 assert src.read().shape == (3, 718, 791) 139 140 141def test_no_initial_bytes(rgb_data_and_profile): 142 """An empty MemoryFile can be opened and written into.""" 143 data, profile = rgb_data_and_profile 144 145 with MemoryFile() as memfile: 146 with memfile.open(**profile) as dst: 147 dst.write(data) 148 view = memfile.getbuffer() 149 # Exact size of the in-memory GeoTIFF varies with GDAL 150 # version and configuration. 151 assert view.size > 1000000 152 # NB: bytes(view) doesn't return what you'd expect with python 2.7. 153 data = bytes(bytearray(view)) 154 155 with MemoryFile(data) as memfile: 156 with memfile.open() as src: 157 assert sorted(src.profile.items()) == sorted(profile.items()) 158 159 160def test_read(tmpdir, rgb_file_bytes): 161 """Reading from a MemoryFile works""" 162 with MemoryFile(rgb_file_bytes) as memfile: 163 tmptiff = tmpdir.join('test.tif') 164 165 while 1: 166 chunk = memfile.read(8192) 167 if not chunk: 168 break 169 tmptiff.write(chunk, 'ab') 170 171 with rasterio.open(str(tmptiff)) as src: 172 assert src.count == 3 173 174 175def test_file_object_read(rgb_file_object): 176 """An example of reading from a file object""" 177 with rasterio.open(rgb_file_object) as src: 178 assert src.driver == 'GTiff' 179 assert src.count == 3 180 assert src.dtypes == ('uint8', 'uint8', 'uint8') 181 assert src.read().shape == (3, 718, 791) 182 183 184def test_file_object_read_variant(rgb_file_bytes): 185 """An example of reading from a MemoryFile object""" 186 with rasterio.open(MemoryFile(rgb_file_bytes)) as src: 187 assert src.driver == 'GTiff' 188 assert src.count == 3 189 assert src.dtypes == ('uint8', 'uint8', 'uint8') 190 assert src.read().shape == (3, 718, 791) 191 192 193def test_file_object_read_variant2(rgb_file_bytes): 194 """An example of reading from a BytesIO object""" 195 with rasterio.open(BytesIO(rgb_file_bytes)) as src: 196 assert src.driver == 'GTiff' 197 assert src.count == 3 198 assert src.dtypes == ('uint8', 'uint8', 'uint8') 199 assert src.read().shape == (3, 718, 791) 200 201 202def test_test_file_object_write(tmpdir, rgb_data_and_profile): 203 """An example of writing to a file object""" 204 data, profile = rgb_data_and_profile 205 with tmpdir.join('test.tif').open('wb') as fout: 206 with rasterio.open(fout, 'w', **profile) as dst: 207 dst.write(data) 208 209 with rasterio.open(str(tmpdir.join('test.tif'))) as src: 210 assert src.driver == 'GTiff' 211 assert src.count == 3 212 assert src.dtypes == ('uint8', 'uint8', 'uint8') 213 assert src.read().shape == (3, 718, 791) 214 215 216def test_nonpersistemt_memfile_fail_example(rgb_data_and_profile): 217 """An example of writing to a file object""" 218 data, profile = rgb_data_and_profile 219 with BytesIO() as fout: 220 with rasterio.open(fout, 'w', **profile) as dst: 221 dst.write(data) 222 223 # This fails because the MemoryFile created in open() is 224 # gone. 225 rasterio.open(fout) 226 227 228def test_zip_closed(): 229 """A closed ZipMemoryFile can not be opened""" 230 with ZipMemoryFile() as zipmemfile: 231 pass 232 with pytest.raises(IOError): 233 zipmemfile.open('foo') 234 235 236def test_zip_file_object_read(path_zip_file): 237 """An example of reading from a zip file object""" 238 with open(path_zip_file, 'rb') as zip_file_object: 239 with ZipMemoryFile(zip_file_object) as zipmemfile: 240 with zipmemfile.open('white-gemini-iv.vrt') as src: 241 assert src.driver == 'VRT' 242 assert src.count == 3 243 assert src.dtypes == ('uint8', 'uint8', 'uint8') 244 assert src.read().shape == (3, 768, 1024) 245 246 247def test_vrt_memfile(): 248 """Successfully read an in-memory VRT""" 249 with open('tests/data/white-gemini-iv.vrt') as vrtfile: 250 source = vrtfile.read() 251 source = source.replace('<SourceFilename relativeToVRT="1">389225main_sw_1965_1024.jpg</SourceFilename>', '<SourceFilename relativeToVRT="0">{}/389225main_sw_1965_1024.jpg</SourceFilename>'.format(os.path.abspath("tests/data"))) 252 253 with MemoryFile(source.encode('utf-8'), ext='vrt') as memfile: 254 with memfile.open() as src: 255 assert src.driver == 'VRT' 256 assert src.count == 3 257 assert src.dtypes == ('uint8', 'uint8', 'uint8') 258 assert src.read().shape == (3, 768, 1024) 259 260 261def test_write_plus_mode(): 262 with MemoryFile() as memfile: 263 with memfile.open(driver='GTiff', dtype='uint8', count=3, height=32, width=32, crs='epsg:3226', transform=Affine.identity() * Affine.scale(0.5, -0.5)) as dst: 264 dst.write(numpy.full((32, 32), 255, dtype='uint8'), 1) 265 dst.write(numpy.full((32, 32), 204, dtype='uint8'), 2) 266 dst.write(numpy.full((32, 32), 153, dtype='uint8'), 3) 267 data = dst.read() 268 assert (data[0] == 255).all() 269 assert (data[1] == 204).all() 270 assert (data[2] == 153).all() 271 272 273def test_write_plus_model_jpeg(): 274 with rasterio.Env(), MemoryFile() as memfile: 275 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: 276 dst.write(numpy.full((32, 32), 255, dtype='uint8'), 1) 277 dst.write(numpy.full((32, 32), 204, dtype='uint8'), 2) 278 dst.write(numpy.full((32, 32), 153, dtype='uint8'), 3) 279 data = dst.read() 280 assert (data[0] == 255).all() 281 assert (data[1] == 204).all() 282 assert (data[2] == 153).all() 283 284 285def test_memfile_copyfiles(path_rgb_msk_byte_tif): 286 """Multiple files can be copied to a MemoryFile using copyfiles""" 287 with rasterio.open(path_rgb_msk_byte_tif) as src: 288 src_basename = os.path.basename(src.name) 289 with MemoryFile(dirname="foo", filename=src_basename) as memfile: 290 copyfiles(src.name, memfile.name) 291 with memfile.open() as rgb2: 292 assert sorted(rgb2.files) == sorted(['/vsimem/foo/{}'.format(src_basename), '/vsimem/foo/{}.msk'.format(src_basename)]) 293 294 295def test_multi_memfile(path_rgb_msk_byte_tif): 296 """Multiple files can be copied to a MemoryFile using copyfiles""" 297 with open(path_rgb_msk_byte_tif, 'rb') as tif_fp: 298 tif_bytes = tif_fp.read() 299 with open(path_rgb_msk_byte_tif + '.msk', 'rb') as msk_fp: 300 msk_bytes = msk_fp.read() 301 302 with MemoryFile(tif_bytes, dirname="bar", filename='foo.tif') as tifmemfile, MemoryFile(msk_bytes, dirname="bar", filename='foo.tif.msk') as mskmemfile: 303 with tifmemfile.open() as src: 304 assert sorted(os.path.basename(fn) for fn in src.files) == sorted(['foo.tif', 'foo.tif.msk']) 305 assert src.mask_flag_enums == ([MaskFlags.per_dataset],) * 3 306 307 308def test_memory_file_gdal_error_message(capsys): 309 """No weird error messages should be seen, see #1659""" 310 memfile = MemoryFile() 311 data = numpy.array([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]]).astype('uint8') 312 west_bound = 0; north_bound = 2; cellsize=0.5; nodata = -9999; driver='AAIGrid'; 313 dtype = data.dtype 314 shape = data.shape 315 transform = rasterio.transform.from_origin(west_bound, north_bound, cellsize, cellsize) 316 dataset = memfile.open(driver=driver, width=shape[1], height=shape[0], transform=transform, count=1, dtype=dtype, nodata=nodata, crs='epsg:3226') 317 dataset.write(data, 1) 318 dataset.close() 319 captured = capsys.readouterr() 320 assert "ERROR 4" not in captured.err 321 assert "ERROR 4" not in captured.out 322 323 324def test_write_plus_mode_requires_width(): 325 """Width is required""" 326 with MemoryFile() as memfile: 327 with pytest.raises(TypeError): 328 memfile.open(driver='GTiff', dtype='uint8', count=3, height=32, crs='epsg:3226', transform=Affine.identity() * Affine.scale(0.5, -0.5)) 329 330 331def test_write_plus_mode_blockxsize_requires_width(): 332 """Width is required""" 333 with MemoryFile() as memfile: 334 with pytest.raises(TypeError): 335 memfile.open(driver='GTiff', dtype='uint8', count=3, height=32, crs='epsg:3226', transform=Affine.identity() * Affine.scale(0.5, -0.5), blockxsize=128) 336 337def test_write_rpcs_to_memfile(): 338 """Ensure we can write rpcs to a new MemoryFile""" 339 with rasterio.open('tests/data/RGB.byte.rpc.vrt') as src: 340 profile = src.profile.copy() 341 with MemoryFile() as memfile: 342 with memfile.open(**profile) as dst: 343 assert dst.rpcs is None 344 dst.rpcs = src.rpcs 345 assert dst.rpcs 346