Improve object handling & testing of ensure_unicode

dask/tests/test_utils.py

@@ -1,8 +1,8 @@
-import array
 import datetime
 import functools
 import operator
 import pickle
+from array import array
 
 import pytest
 from tlz import curry
@@ -21,6 +21,7 @@
     ensure_bytes,
     ensure_dict,
     ensure_set,
+    ensure_unicode,
     extra_titles,
     factors,
     format_bytes,
@@ -47,7 +48,7 @@
 
 
 def test_ensure_bytes():
-    data = [b"1", "1", memoryview(b"1"), bytearray(b"1"), array.array("b", [49])]
+    data = [b"1", "1", memoryview(b"1"), bytearray(b"1"), array("B", b"1")]
     for d in data:
         result = ensure_bytes(d)
         assert isinstance(result, bytes)
@@ -67,6 +68,30 @@ def test_ensure_bytes_pyarrow_buffer():
     assert isinstance(result, bytes)
 
 
+def test_ensure_unicode():
+    data = [b"1", "1", memoryview(b"1"), bytearray(b"1"), array("B", b"1")]
+    for d in data:
+        result = ensure_unicode(d)
+        assert isinstance(result, str)
+        assert result == "1"
+
+
+def test_ensure_unicode_ndarray():
+    np = pytest.importorskip("numpy")
+    a = np.frombuffer(b"123", dtype="u1")
+    result = ensure_unicode(a)
+    assert isinstance(result, str)
+    assert result == "123"
+
+
+def test_ensure_unicode_pyarrow_buffer():
+    pa = pytest.importorskip("pyarrow")
+    buf = pa.py_buffer(b"123")
+    result = ensure_unicode(buf)
+    assert isinstance(result, str)
+    assert result == "123"
+
+
 def test_getargspec():
     def func(x, y):
         pass

dask/utils.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import codecs
 import functools
 import inspect
 import os
@@ -921,7 +922,7 @@ def ensure_bytes(s) -> bytes:
             return bytes(s)
         except Exception as e:
             raise TypeError(
-                f"Object {s} is neither a bytes object nor has an encode method"
+                f"Object {s} is neither a bytes object nor can be encoded to bytes"
             ) from e
 
 
@@ -935,9 +936,15 @@ def ensure_unicode(s) -> str:
     """
     if isinstance(s, str):
         return s
-    if hasattr(s, "decode"):
+    elif hasattr(s, "decode"):
         return s.decode()
-    raise TypeError(f"Object {s} is neither a str object nor has an decode method")
+    else:
+        try:
+            return codecs.decode(s)
+        except Exception as e:
+            raise TypeError(
+                f"Object {s} is neither a str object nor can be decoded to str"
+            ) from e
 
 
 def digit(n, k, base):