xarray.backends refactor

xarray/backends/common.py

@@ -508,30 +508,3 @@ 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, file, mode='r', **kwargs):
-        self.opener = partial(opener, file, mode=mode, **kwargs)
-        self.mode = mode
-        self._ds = None
-
-    @property
-    def value(self):
-        self._ds = self.opener()
-        return self._ds
-
-    def __getstate__(self):
-        state = self.__dict__.copy()
-        del state['_ds']
-        if self.mode == 'w':
-            # file has already been created, don't override when restoring
-            state['mode'] = 'a'
-        return state
-
-    def __setstate__(self, state):
-        self.__dict__.update(state)
-
-    def close(self):
-        self._ds.close()

xarray/backends/file_manager.py

@@ -0,0 +1,118 @@
+import contextlib
+import threading
+
+from .lru_cache import LRUCache
+
+
+# Global cache for storing open files.
+FILE_CACHE = LRUCache(512, on_evict=lambda k, v: v.close())
+
+# TODO(shoyer): add an option (xarray.set_options) for resizing the cache.
+
+
+class FileManager(object):
+    """Wrapper for automatically opening and closing file objects.
+
+    Unlike files, FileManager objects can be safely pickled and passed between
+    processes. They should be explicitly closed to release resources, but
+    a per-process least-recently-used cache for open files ensures that you can
+    safely create arbitrarily large numbers of FileManager objects.
+
+    Example usage:
+
+        manager = FileManager(open, 'example.txt', mode='w')
+        with manager.acquire() as f:
+            f.write(...)
+        manager.close()
+    """
+
+    def __init__(self, opener, *args, **kwargs):
+        """Initialize a FileManager.
+
+        Parameters
+        ----------
+        opener : callable
+            Function that when called like ``opener(*args, **kwargs)`` returns
+            an open file object. The file object must implement a ``close()``
+            method.
+        *args
+            Positional arguments for opener. A ``mode`` argument should be
+            provided as a keyword argument (see below).
+        **kwargs
+            Keyword arguments for opener. The keyword argument ``mode`` has
+            special handling if it is provided with a value of 'w': on all
+            calls after the first, it is changed to 'a' instead to avoid
+            overriding the newly created file. All argument values must be
+            hashable.
+        """
+        self._opener = opener
+        self._args = args
+        self._kwargs = kwargs
+        self._key = _make_key(opener, args, kwargs)
+        self._lock = threading.RLock()
+
+    @contextlib.contextmanager
+    def acquire(self):
+        """Context manager for acquiring a file object.
+
+        A new file is only opened if it has expired from the
+        least-recently-used cache.
+
+        This method uses a reentrant lock, which ensures that it is
+        thread-safe. You can safely acquire a file in multiple threads at the
+        same time, as long as the underlying file object is thread-safe.
+
+        Yields
+        ------
+        Open file object, as returned by ``opener(*args, **kwargs)``.
+        """
+        with self._lock:
+            try:
+                file = FILE_CACHE[self._key]
+            except KeyError:
+                file = self._opener(*self._args, **self._kwargs)
+                if self._kwargs.get('mode') == 'w':
+                    # ensure file doesn't get overriden when opened again
+                    self._kwargs['mode'] = 'a'
+                    self._key = _make_key(
+                        self._opener, self._args, self._kwargs)
+                FILE_CACHE[self._key] = file
+        yield file
+
+    def close(self):
+        """Explicitly close any associated file object (if necessary)."""
+        file = FILE_CACHE.pop(self._key, default=None)
+        if file is not None:
+            file.close()
+
+    def __getstate__(self):
+        """State for pickling."""
+        return (self._opener, self._args, self._kwargs)
+
+    def __setstate__(self, state):
+        """Restore from a pickle."""
+        opener, args, kwargs = state
+        self.__init__(opener, *args, **kwargs)
+
+
+class _HashedSequence(list):
+    """Speedup repeated look-ups by caching hash values.
+
+    Based on what Python uses internally in functools.lru_cache.
+
+    Python doesn't perform this optimization automatically:
+    https://bugs.python.org/issue1462796
+    """
+
+    def __init__(self, tuple_value):
+        self[:] = tuple_value
+        self.hashvalue = hash(tuple_value)
+
+    def __hash__(self):
+        return self.hashvalue
+
+
+def _make_key(opener, args, kwargs):
+    """Make a key for caching files in the LRU cache."""
+    value = (opener, args, tuple(sorted(kwargs.items())))
+    return _HashedSequence(value)

xarray/backends/lru_cache.py

@@ -0,0 +1,78 @@
+import collections
+import threading
+
+
+class LRUCache(collections.MutableMapping):
+    """Thread-safe LRUCache based on an OrderedDict.
+
+    All dict operations (__getitem__, __setitem__, __contains__) update the
+    priority of the relevant key and take O(1) time. The dict is iterated over
+    in order from the oldest to newest key, which means that a complete pass
+    over the dict should not affect the order of any entries.
+
+    When a new item is set and the maximum size of the cache is exceeded, the
+    oldest item is dropped and called with ``on_evict(key, value)``.
+
+    The ``maxsize`` property can be used to view or resize the capacity of
+    the cache.
+    """
+    def __init__(self, maxsize, on_evict=None):
+        """
+        Parameters
+        ----------
+        maxsize : int
+            Integer maximum number of items to hold in the cache.
+        on_evict: callable, optional
+            Function to call like ``on_evict(key, value)`` when items are
+            evicted.
+        """
+        self._maxsize = maxsize
+        self._on_evict = on_evict
+        self._cache = collections.OrderedDict()
+        self._lock = threading.RLock()
+
+    def __getitem__(self, key):
+        # record recent use of the key by moving it to the front of the list
+        with self._lock:
+            value = self._cache[key]
+            self._cache.move_to_end(key)
+            return value
+
+    def _shrink(self, capacity):
+        """Shrink the cache if necessary, evicting the oldest items."""
+        while len(self._cache) > capacity:
+            key, value = self._cache.popitem(last=False)
+            if self._on_evict is not None:
+                self._on_evict(key, value)
+
+    def __setitem__(self, key, value):
+        with self._lock:
+            if self._maxsize:
+                if key in self._cache:
+                    self._cache.move_to_end(key)
+                elif len(self._cache) >= self._maxsize:
+                    self._shrink(self._maxsize - 1)
+                self._cache[key] = value
+
+    def __delitem__(self, key):
+        del self._cache[key]
+
+    def __iter__(self):
+        # create a list, so accessing the cache during iteration cannot change
+        # the iteration order
+        return iter(list(self._cache))
+
+    def __len__(self):
+        return len(self._cache)
+
+    @property
+    def maxsize(self):
+        """Maximum number of items can be held in the cache."""
+        return self._maxsize
+
+    @maxsize.setter
+    def maxsize(self, size):
+        """Resize the cache, evicting the oldest items if necessary."""
+        with self._lock:
+            self._shrink(size)
+            self._maxsize = size