Feature/pickle rasterio

doc/whats-new.rst

@@ -86,6 +86,9 @@ Bug fixes
   By `Deepak Cherian <https://github.com/dcherian>`_.
 - Colorbar limits are now determined by excluding ±Infs too.
   By `Deepak Cherian <https://github.com/dcherian>`_.
+- Fixed a bug in ``rasterio`` backend which prevented use with ``distributed``.
+  The ``rasterio`` backend now returns pickleable objects (:issue:`2021`).
+  By `Joe Hamman <https://github.com/jhamman>`_.
 - Fixed ``to_iris`` to maintain lazy dask array after conversion (:issue:`2046`).
   By `Alex Hilson <https://github.com/AlexHilson>`_ and `Stephan Hoyer <https://github.com/shoyer>`_.
 

xarray/backends/common.py

@@ -503,3 +503,43 @@ def assert_open(self):
         if not self._isopen:
             raise AssertionError('internal failure: file must be open '
                                  'if `autoclose=True` is used.')
+
+
+class PickleByReconstructionWrapper(object):
+
+    def __init__(self, opener, *args, **kwargs):
+
+        self.opener = opener
+        self.open_args = args
+        self.open_kwargs = kwargs
+
+        self._ds = None
+        self._isopen = False
+
+    @property
+    def value(self):
+        if self._ds is not None and self._isopen:
+            return self._ds
+        self._ds = self.opener(*self.open_args, **self.open_kwargs)
+        self._isopen = True
+        return self._ds
+
+    def __getstate__(self):
+        state = self.__dict__.copy()
+        del state['_ds']
+        del state['_isopen']
+        if self.open_kwargs.get('mode', None) == 'w':
+            # file has already been created, don't override when restoring
+            state['open_kwargs']['mode'] = 'a'
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._ds = None
+        self._isopen = False
+
+    def __getitem__(self, key):
+        return self.value[key]
+
+    def __getattr__(self, name):
+        return getattr(self.value, name)

xarray/backends/rasterio_.py

@@ -8,7 +8,7 @@
 from .. import DataArray
 from ..core import indexing
 from ..core.utils import is_scalar
-from .common import BackendArray
+from .common import BackendArray, PickleByReconstructionWrapper
 
 try:
     from dask.utils import SerializableLock as Lock
@@ -194,7 +194,8 @@ def open_rasterio(filename, parse_coordinates=None, chunks=None, cache=None,
     """
 
     import rasterio
-    riods = rasterio.open(filename, mode='r')
+
+    riods = PickleByReconstructionWrapper(rasterio.open, filename, mode='r')
 
     if cache is None:
         cache = chunks is None

xarray/tests/test_backends.py

@@ -2649,6 +2649,14 @@ def test_chunks(self):
                 ex = expected.sel(band=1).mean(dim='x')
                 assert_allclose(ac, ex)
 
+    def test_pickle_rio(self):
+        # regression test for https://github.com/pydata/xarray/issues/2121
+        with create_tmp_geotiff() as (tmp_file, expected):
+            with xr.open_rasterio(tmp_file) as rioda:
+                temp = pickle.dumps(rioda)
+                with pickle.loads(temp) as actual:
+                    assert_equal(actual, rioda)
+
     def test_ENVI_tags(self):
         rasterio = pytest.importorskip('rasterio', minversion='1.0a')
         from rasterio.transform import from_origin