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