Merge pull request #30 from dask/mapping

Add mutable mapping

docs/source/api.rst

@@ -31,8 +31,19 @@ API
    S3File.tell
    S3File.write
 
+.. currentmodule:: s3fs.mapping
+
+.. autosummary::
+   S3Map
+
+.. currentmodule:: s3fs.core
+
 .. autoclass:: S3FileSystem
    :members:
 
 .. autoclass:: S3File
    :members:
+
+.. currentmodule:: s3fs.mapping
+
+.. autoclass:: S3Map

s3fs/mapping.py

@@ -0,0 +1,78 @@
+
+from collections import MutableMapping
+import os
+
+class S3Map(MutableMapping):
+    """Wrap an S3FileSystem as a mutable wrapping.
+
+    The keys of the mapping become files under the given root, and the
+    values (which must be bytes) the contents of those files.
+
+    Parameters
+    ----------
+    s3 : S3FileSystem
+    root : string
+        prefix for all the files (perhaps justa  bucket name
+    check : bool (=True)
+        performs a touch at the location, to check writeability.
+
+    Examples
+    --------
+    >>> s3 = s3fs.S3FileSystem() # doctest: +SKIP
+    >>> mw = MapWrapping(s3, 'mybucket/mapstore/') # doctest: +SKIP
+    >>> mw['loc1'] = b'Hello World' # doctest: +SKIP
+    >>> list(mw.keys()) # doctest: +SKIP
+    ['loc1']
+    >>> mw['loc1'] # doctest: +SKIP
+    b'Hello World'
+    """
+
+    def __init__(self, s3, root, check=False):
+        self.s3 = s3
+        self.root = root
+        if check:
+            s3.touch(root+'/a')
+            s3.rm(root+'/a')
+
+    def clear(self):
+        """Remove all keys below root - empties out mapping
+        """
+        self.s3.rm(self.root, recursive=True)
+
+    def _key_to_str(self, key):
+        if isinstance(key, (tuple, list)):
+            key = str(tuple(key))
+        else:
+            key = str(key)
+        return '/'.join([self.root, key])
+
+    def __getitem__(self, key):
+        key = self._key_to_str(key)
+        try:
+            with self.s3.open(key, 'rb') as f:
+                result = f.read()
+        except (IOError, OSError):
+            raise KeyError(key)
+        return result
+
+    def __setitem__(self, key, value):
+        key = self._key_to_str(key)
+        if not isinstance(value, bytes):
+            raise TypeError("Value must be of type bytes")
+        with self.s3.open(key, 'wb') as f:
+            f.write(value)
+
+    def keys(self):
+        return (x[len(self.root) + 1:] for x in self.s3.walk(self.root))
+
+    def __iter__(self):
+        return self.keys()
+
+    def __delitem__(self, key):
+        self.s3.rm(self._key_to_str(key))
+
+    def __contains__(self, key):
+        return self.s3.exists(self._key_to_str(key))
+
+    def __len__(self):
+        return sum(1 for _ in self.keys())

s3fs/tests/test_mapping.py

@@ -0,0 +1,51 @@
+from s3fs.tests.test_s3fs import s3, test_bucket_name
+from s3fs.mapping import S3Map
+
+root = test_bucket_name+'/mapping'
+
+
+def test_simple(s3):
+    mw = S3Map(s3, root)
+    assert not mw
+
+    assert list(mw) == list(mw.keys()) == []
+    assert list(mw.values()) == []
+    assert list(mw.items()) == []
+
+
+def test_with_data(s3):
+    mw = S3Map(s3, root)
+    mw['x'] = b'123'
+    assert list(mw) == list(mw.keys()) == ['x']
+    assert list(mw.values()) == [b'123']
+    assert list(mw.items()) == [('x', b'123')]
+    assert mw['x'] == b'123'
+    assert bool(mw)
+
+    assert s3.walk(root) == [test_bucket_name+'/mapping/x']
+    mw['x'] = b'000'
+    assert mw['x'] == b'000'
+
+    mw['y'] = b'456'
+    assert mw['y'] == b'456'
+    assert set(mw) == {'x', 'y'}
+
+    mw.clear()
+    assert list(mw) == []
+
+
+def test_complex_keys(s3):
+    mw = S3Map(s3, root)
+    mw[1] = b'hello'
+    assert mw[1] == b'hello'
+    del mw[1]
+
+    mw[1, 2] = b'world'
+    assert mw[1, 2] == b'world'
+    del mw[1, 2]
+
+    mw['x', 1, 2] = b'hello world'
+    assert mw['x', 1, 2] == b'hello world'
+    print(list(mw))
+
+    assert ('x', 1, 2) in mw