Repository: rasterio/rasterio

Issue:
Lightweight context managers around file objects and memory files
It took me long enough, but now I get what `contextlib.contextmanager` is about :)

Resolves #964


Discussion:

Resolving patch:
diff --git a/rasterio/__init__.py b/rasterio/__init__.py
--- a/rasterio/__init__.py
+++ b/rasterio/__init__.py
@@ -66,25 +66,33 @@ def window_index(*args, **kwargs):
 log.addHandler(NullHandler())
 
 
-def open(fp, mode='r', driver=None, width=None, height=None,
-         count=None, crs=None, transform=None, dtype=None, nodata=None,
-         **kwargs):
-    """Open file at ``path`` in ``mode`` 'r' (read), 'r+' (read and
-    write), or 'w' (write) and return a dataset Reader or Updater
-    object.
+def open(fp, mode='r', driver=None, width=None, height=None, count=None,
+         crs=None, transform=None, dtype=None, nodata=None, **kwargs):
+    """Open a dataset for reading or writing.
+
+    The dataset may be located in a local file, in a resource located
+    by a URL, or contained within a stream of bytes.
+
+    To access a dataset within a zip file without unzipping the archive
+    use an Apache VFS style zip:// URL like
+
+      zip://path/to/archive.zip!path/to/example.tif
+
+    In read ('r') or read/write ('r+') mode, no other keyword arguments
+    are required: the attributes are supplied by the opened dataset.
 
     In write mode, a driver name such as "GTiff" or "JPEG" (see GDAL
-    docs or ``gdal_translate --help`` on the command line),
-    ``width`` (number of pixels per line) and ``height`` (number of
-    lines), the ``count`` number of bands in the new file must be
-    specified.  Additionally, the data type for bands such as
-    ``rasterio.ubyte`` for 8-bit bands or ``rasterio.uint16`` for
-    16-bit bands must be specified using the ``dtype`` argument.
+    docs or ``gdal_translate --help`` on the command line), ``width``
+    (number of pixels per line) and ``height`` (number of lines), the
+    ``count`` number of bands in the new file must be specified.
+    Additionally, the data type for bands such as ``rasterio.ubyte`` for
+    8-bit bands or ``rasterio.uint16`` for 16-bit bands must be
+    specified using the ``dtype`` argument.
 
     Parameters
     ----------
     fp: string or file
-        A filename or file object opened in binary mode.
+        A filename or URL, or file object opened in binary mode.
     mode: string
         "r" (read), "r+" (read/write), or "w" (write)
     driver: string
@@ -240,7 +248,7 @@ def fp_writer(fp):
         return fp_writer(fp)
 
     else:
-        # The 'normal' filename path.
+        # The 'normal' filename or URL path.
         _, _, scheme = parse_path(fp)
 
         with Env() as env:
diff --git a/rasterio/io.py b/rasterio/io.py
--- a/rasterio/io.py
+++ b/rasterio/io.py
@@ -218,8 +218,7 @@ class MemoryFile(MemoryFileBase):
     A GeoTIFF can be loaded in memory and accessed using the GeoTIFF
     format driver
 
-    >>> with open('tests/data/RGB.byte.tif', 'rb') as f, \
-    ...         MemoryFile(f.read()) as memfile:
+    >>> with open('tests/data/RGB.byte.tif', 'rb') as f, MemoryFile(f) as memfile:
     ...     with memfile.open() as src:
     ...         pprint.pprint(src.profile)
     ...
@@ -236,26 +235,37 @@ class MemoryFile(MemoryFileBase):
      'width': 791}
 
     """
+    def __init__(self, file_or_bytes=None, ext=''):
+        super(MemoryFile, self).__init__(file_or_bytes=file_or_bytes, ext=ext)
 
-    def open(self, driver=None, width=None, height=None,
-             count=None, crs=None, transform=None, dtype=None, nodata=None,
-             **kwargs):
+    def open(self, driver=None, width=None, height=None, count=None, crs=None,
+             transform=None, dtype=None, nodata=None, **kwargs):
         """Open the file and return a Rasterio dataset object.
 
         If data has already been written, the file is opened in 'r+'
         mode. Otherwise, the file is opened in 'w' mode.
+
+        Parameters
+        ----------
+        Note well that there is no `path` parameter: a `MemoryFile`
+        contains a single dataset and there is no need to specify a
+        path.
+
+        Other parameters are optional and have the same semantics as the
+        parameters of `rasterio.open()`.
         """
+        vsi_path = self.name
+
         with Env():
             if self.closed:
                 raise IOError("I/O operation on closed file.")
             if self.exists():
-                s = get_writer_for_path(self.name)(self.name, 'r+')
+                s = DatasetReader(vsi_path, 'r+')
             else:
-                s = get_writer_for_driver(
-                    driver)(self.name, 'w', driver=driver, width=width,
-                            height=height, count=count, crs=crs,
-                            transform=transform, dtype=dtype, nodata=nodata,
-                            **kwargs)
+                s = DatasetWriter(vsi_path, 'w', driver=driver, width=width,
+                                  height=height, count=count, crs=crs,
+                                  transform=transform, dtype=dtype,
+                                  nodata=nodata, **kwargs)
             s.start()
             return s
 
@@ -266,6 +276,39 @@ def __exit__(self, *args, **kwargs):
         self.close()
 
 
+class ZipMemoryFile(MemoryFile):
+    """A read-only BytesIO-like object backed by an in-memory zip file.
+
+    This allows a zip file containing formatted files to be read
+    without I/O.
+    """
+
+    def __init__(self, file_or_bytes=None):
+        super(ZipMemoryFile, self).__init__(file_or_bytes, ext='zip')
+
+    def open(self, path):
+        """Open a dataset within the zipped stream.
+
+        Parameters
+        ----------
+        path : str
+            Path to a dataset in the zip file, relative to the root of the
+            archive.
+
+        Returns
+        -------
+        A Rasterio dataset object
+        """
+        vsi_path = '/vsizip{0}/{1}'.format(self.name, path.lstrip('/'))
+
+        with Env():
+            if self.closed:
+                raise IOError("I/O operation on closed file.")
+            s = DatasetReader(vsi_path, 'r')
+            s.start()
+            return s
+
+
 def get_writer_for_driver(driver):
     """Return the writer class appropriate for the specified driver."""
     cls = None
diff --git a/rasterio/vfs.py b/rasterio/vfs.py
--- a/rasterio/vfs.py
+++ b/rasterio/vfs.py
@@ -59,9 +59,10 @@ def vsi_path(path, archive=None, scheme=None):
     elif scheme and scheme == 's3':
         result = "/vsis3/{0}".format(path)
     elif scheme and scheme != 'file':
-        path = path.strip('/')
-        result = '/'.join(
-            ['/vsi{0}'.format(scheme), archive, path])
+        if archive:
+            result = '/vsi{0}/{1}/{2}'.format(scheme, archive, path.lstrip('/'))
+        else:
+            result = '/vsi{0}/{1}'.format(scheme, path.lstrip('/'))
     else:
         result = path
     return result


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,6 +3,7 @@
 import os
 import shutil
 import sys
+import zipfile
 
 from click.testing import CliRunner
 import py
@@ -252,6 +253,25 @@ def fin():
     request.addfinalizer(fin)
 
 
-@pytest.fixture(scope='module')
-def path_rgb_byte_tif():
-    return os.path.join('tests', 'data', 'RGB.byte.tif')
+@pytest.fixture(scope='session')
+def data_dir():
+    """Absolute file path to the directory containing test datasets."""
+    return os.path.abspath(os.path.join('tests', 'data'))
+
+
+@pytest.fixture(scope='session')
+def path_rgb_byte_tif(data_dir):
+    return os.path.join(data_dir, 'RGB.byte.tif')
+
+
+@pytest.fixture(scope='session')
+def path_zip_file():
+    """Creates ``coutwildrnp.zip`` if it does not exist and returns
+    the absolute file path."""
+    path = '{}/white-gemini-iv.zip'.format(data_dir())
+    if not os.path.exists(path):
+        with zipfile.ZipFile(path, 'w') as zip:
+            for filename in ['white-gemini-iv.vrt',
+                             '389225main_sw_1965_1024.jpg']:
+                zip.write(os.path.join(data_dir(), filename), filename)
+    return path
diff --git a/tests/test_memoryfile.py b/tests/test_memoryfile.py
--- a/tests/test_memoryfile.py
+++ b/tests/test_memoryfile.py
@@ -1,12 +1,13 @@
 """MemoryFile tests"""
 
+from io import BytesIO
 import logging
 
 from packaging.version import parse
 import pytest
 
 import rasterio
-from rasterio.io import MemoryFile
+from rasterio.io import MemoryFile, ZipMemoryFile
 
 
 logging.basicConfig(level=logging.DEBUG)
@@ -18,16 +19,33 @@
     reason="MemoryFile requires GDAL 2.0")
 
 
-@pytest.fixture(scope='function')
+@pytest.fixture(scope='session')
 def rgb_file_bytes(path_rgb_byte_tif):
+    """Get the bytes of our RGB.bytes.tif file"""
     return open(path_rgb_byte_tif, 'rb').read()
 
 
 @pytest.fixture(scope='function')
 def rgb_file_object(path_rgb_byte_tif):
+    """Get RGB.bytes.tif file opened in 'rb' mode"""
     return open(path_rgb_byte_tif, 'rb')
 
 
+@pytest.fixture(scope='session')
+def rgb_data_and_profile(path_rgb_byte_tif):
+    with rasterio.open(path_rgb_byte_tif) as src:
+        data = src.read()
+        profile = src.profile
+    return data, profile
+
+
+@mingdalversion
+def test_initial_not_bytes():
+    """Creating a MemoryFile from not bytes fails."""
+    with pytest.raises(TypeError):
+        MemoryFile(u'lolwut')
+
+
 @mingdalversion
 def test_initial_bytes(rgb_file_bytes):
     """MemoryFile contents can initialized from bytes and opened."""
@@ -39,6 +57,26 @@ def test_initial_bytes(rgb_file_bytes):
             assert src.read().shape == (3, 718, 791)
 
 
+@mingdalversion
+def test_initial_file_object(rgb_file_object):
+    """MemoryFile contents can initialized from bytes and opened."""
+    with MemoryFile(rgb_file_object) as memfile:
+        with memfile.open() as src:
+            assert src.driver == 'GTiff'
+            assert src.count == 3
+            assert src.dtypes == ('uint8', 'uint8', 'uint8')
+            assert src.read().shape == (3, 718, 791)
+
+
+@mingdalversion
+def test_closed():
+    """A closed MemoryFile can not be opened"""
+    with MemoryFile() as memfile:
+        pass
+    with pytest.raises(IOError):
+        memfile.open()
+
+
 @mingdalversion
 def test_non_initial_bytes(rgb_file_bytes):
     """MemoryFile contents can be read from bytes and opened."""
@@ -76,14 +114,6 @@ def test_non_initial_bytearray(rgb_file_bytes):
             assert src.read().shape == (3, 718, 791)
 
 
-@pytest.fixture(scope='function')
-def rgb_data_and_profile(path_rgb_byte_tif):
-    with rasterio.open(path_rgb_byte_tif) as src:
-        data = src.read()
-        profile = src.profile
-    return data, profile
-
-
 @mingdalversion
 def test_no_initial_bytes(rgb_data_and_profile):
     """An empty MemoryFile can be opened and written into."""
@@ -129,6 +159,26 @@ def test_file_object_read(rgb_file_object):
         assert src.read().shape == (3, 718, 791)
 
 
+@mingdalversion
+def test_file_object_read_variant(rgb_file_bytes):
+    """An example of reading from a MemoryFile object"""
+    with rasterio.open(MemoryFile(rgb_file_bytes)) as src:
+        assert src.driver == 'GTiff'
+        assert src.count == 3
+        assert src.dtypes == ('uint8', 'uint8', 'uint8')
+        assert src.read().shape == (3, 718, 791)
+
+
+@mingdalversion
+def test_file_object_read_variant2(rgb_file_bytes):
+    """An example of reading from a BytesIO object"""
+    with rasterio.open(BytesIO(rgb_file_bytes)) as src:
+        assert src.driver == 'GTiff'
+        assert src.count == 3
+        assert src.dtypes == ('uint8', 'uint8', 'uint8')
+        assert src.read().shape == (3, 718, 791)
+
+
 @mingdalversion
 def test_test_file_object_write(tmpdir, rgb_data_and_profile):
     """An example of writing to a file object"""
@@ -142,3 +192,39 @@ def test_test_file_object_write(tmpdir, rgb_data_and_profile):
         assert src.count == 3
         assert src.dtypes == ('uint8', 'uint8', 'uint8')
         assert src.read().shape == (3, 718, 791)
+
+
+@mingdalversion
+@pytest.mark.xfail(reason="in-memory files created within open() "
+                          "are not persistent")
+def test_nonpersistemt_memfile_fail_example(rgb_data_and_profile):
+    """An example of writing to a file object"""
+    data, profile = rgb_data_and_profile
+    with BytesIO() as fout:
+        with rasterio.open(fout, 'w', **profile) as dst:
+            dst.write(data)
+
+        # This fails because the MemoryFile created in open() is
+        # gone.
+        rasterio.open(fout)
+
+
+@mingdalversion
+def test_zip_closed():
+    """A closed ZipMemoryFile can not be opened"""
+    with ZipMemoryFile() as zipmemfile:
+        pass
+    with pytest.raises(IOError):
+        zipmemfile.open('foo')
+
+
+@mingdalversion
+def test_zip_file_object_read(path_zip_file):
+    """An example of reading from a zip file object"""
+    with open(path_zip_file, 'rb') as zip_file_object:
+        with ZipMemoryFile(zip_file_object) as zipmemfile:
+            with zipmemfile.open('white-gemini-iv.vrt') as src:
+                assert src.driver == 'VRT'
+                assert src.count == 3
+                assert src.dtypes == ('uint8', 'uint8', 'uint8')
+                assert src.read().shape == (3, 768, 1024)