Improve object handling & testing of ensure_unicode

dask/tests/test_utils.py

@@ -21,6 +21,7 @@
     ensure_bytes,
     ensure_dict,
     ensure_set,
+    ensure_unicode,
     extra_titles,
     factors,
     format_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.array("b", [49])]
+    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
@@ -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 has an decode method"
+            ) from e
 
 
 def digit(n, k, base):