From dcd23a9fe1562814e7682695e635e6d42ff142bd Mon Sep 17 00:00:00 2001
From: Martin Durant <martin.durant@utoronto.ca>
Date: Thu, 19 Nov 2020 09:59:04 -0500
Subject: [PATCH] Allow various bytelike to mapper setitem(s)

---
 fsspec/mapping.py            | 14 +++++++++-----
 fsspec/tests/test_mapping.py | 22 ++++++++++++++++++++++
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/fsspec/mapping.py b/fsspec/mapping.py
index 5cb54ee1a..d453d96ea 100644
--- a/fsspec/mapping.py
+++ b/fsspec/mapping.py
@@ -106,7 +106,7 @@ def setitems(self, values_dict):
         ----------
         values_dict: dict(str, bytes)
         """
-        values = {self._key_to_str(k): v for k, v in values_dict.items()}
+        values = {self._key_to_str(k): maybe_convert(v) for k, v in values_dict.items()}
         self.fs.pipe(values)
 
     def delitems(self, keys):
@@ -147,11 +147,8 @@ def pop(self, key, default=None):
     def __setitem__(self, key, value):
         """Store value in key"""
         key = self._key_to_str(key)
-        if isinstance(value, array.array):  # pragma: no cover
-            # back compat, array.array used to work
-            value = bytearray(value)
         self.fs.mkdirs(self.fs._parent(key), exist_ok=True)
-        self.fs.pipe_file(key, value)
+        self.fs.pipe_file(key, maybe_convert(value))
 
     def __iter__(self):
         return (self._str_to_key(x) for x in self.fs.find(self.root))
@@ -175,6 +172,13 @@ def __reduce__(self):
         return FSMap, (self.root, self.fs, False, False, self.missing_exceptions)
 
 
+def maybe_convert(value):
+    if isinstance(value, array.array) or hasattr(value, "__array__"):
+        # bytes-like things
+        value = bytearray(memoryview(value))
+    return value
+
+
 def get_mapper(url, check=False, create=False, missing_exceptions=None, **kwargs):
     """Create key-value interface for given URL and options
 
diff --git a/fsspec/tests/test_mapping.py b/fsspec/tests/test_mapping.py
index 31102eba5..5992de74e 100644
--- a/fsspec/tests/test_mapping.py
+++ b/fsspec/tests/test_mapping.py
@@ -88,3 +88,25 @@ def test_multi():
     assert m.getitems(list(data)) == data
     m.delitems(list(data))
     assert not list(m)
+
+
+def test_setitem_types():
+    import array
+
+    m = fsspec.get_mapper("memory://")
+    m["a"] = array.array("i", [1])
+    assert m["a"] == b"\x01\x00\x00\x00"
+    m["b"] = bytearray(b"123")
+    assert m["b"] == b"123"
+    m.setitems({"c": array.array("i", [1]), "d": bytearray(b"123")})
+    assert m["c"] == b"\x01\x00\x00\x00"
+    assert m["d"] == b"123"
+
+
+def test_setitem_numpy():
+    m = fsspec.get_mapper("memory://")
+    np = pytest.importorskip("numpy")
+    m["c"] = np.array(1, dtype="int32")  # scalar
+    assert m["c"] == b"\x01\x00\x00\x00"
+    m["c"] = np.array([1, 2], dtype="int32")  # array
+    assert m["c"] == b"\x01\x00\x00\x00\x02\x00\x00\x00"