Merge pull request #39 from dask/bytearray

Ensure s3fs works with byte arrays

s3fs/core.py

@@ -12,7 +12,7 @@
 from botocore.exceptions import ClientError, ParamValidationError
 from botocore.client import Config
 
-from .utils import read_block, raises
+from .utils import read_block, raises, ensure_writable
 
 logger = logging.getLogger(__name__)
 
@@ -783,7 +783,7 @@ def write(self, data):
             raise ValueError('File not in write mode')
         if self.closed:
             raise ValueError('I/O operation on closed file.')
-        out = self.buffer.write(data)
+        out = self.buffer.write(ensure_writable(data))
         self.loc += out
         if self.buffer.tell() > self.blocksize:
             self.flush()

s3fs/mapping.py

@@ -64,8 +64,6 @@ def __getitem__(self, key):
 
     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)
 

s3fs/tests/test_mapping.py

@@ -75,3 +75,19 @@ def test_pickle(s3):
     d2 = pickle.loads(pickle.dumps(d))
 
     assert d2['x'] == b'1'
+
+
+def test_array(s3):
+    from array import array
+    d = S3Map(root, s3)
+    d['x'] = array('B', [65] * 1000)
+
+    assert d['x'] == b'A' * 1000
+
+
+def test_bytearray(s3):
+    from array import array
+    d = S3Map(root, s3)
+    d['x'] = bytearray(b'123')
+
+    assert d['x'] == b'123'

s3fs/tests/test_s3fs.py

@@ -629,3 +629,15 @@ def test_bigger_than_block_read(s3):
 
 def test_current(s3):
     assert S3FileSystem.current() is s3
+
+
+def test_array(s3):
+    from array import array
+    data = array('B', [65] * 1000)
+
+    with s3.open(a, 'wb') as f:
+        f.write(data)
+
+    with s3.open(a, 'rb') as f:
+        out = f.read()
+        assert out == b'A' * 1000

s3fs/utils.py

@@ -1,8 +1,13 @@
 
+import array
 from contextlib import contextmanager
 import os
 import tempfile
 import shutil
+import sys
+
+PY2 = sys.version_info[0] == 2
+PY3 = sys.version_info[0] == 3
 
 
 @contextmanager
@@ -124,3 +129,9 @@ def raises(exc, lamda):
         return False
     except exc:
         return True
+
+
+def ensure_writable(b):
+    if PY2 and isinstance(b, array.array):
+        return b.tostring()
+    return b