Use ensure_bytes from dask.utils

distributed/comm/ws.py

@@ -15,6 +15,7 @@
 from tornado.websocket import WebSocketClosedError, WebSocketHandler, websocket_connect
 
 import dask
+from dask.utils import ensure_bytes
 
 from distributed.comm.addressing import parse_host_port, unparse_host_port
 from distributed.comm.core import (
@@ -36,7 +37,7 @@
     get_tcp_server_address,
     to_frames,
 )
-from distributed.utils import ensure_bytes, nbytes
+from distributed.utils import nbytes
 
 logger = logging.getLogger(__name__)
 

distributed/protocol/compression.py

@@ -15,8 +15,7 @@
 from tlz import identity
 
 import dask
-
-from distributed.utils import ensure_bytes
+from dask.utils import ensure_bytes
 
 compressions: dict[
     str | None | Literal[False],

distributed/protocol/serialize.py

@@ -11,7 +11,7 @@
 
 import dask
 from dask.base import normalize_token
-from dask.utils import typename
+from dask.utils import ensure_bytes, typename
 
 from distributed.protocol import pickle
 from distributed.protocol.compression import decompress, maybe_compress
@@ -22,7 +22,7 @@
     pack_frames_prelude,
     unpack_frames,
 )
-from distributed.utils import ensure_bytes, has_keyword
+from distributed.utils import has_keyword
 
 dask_serialize = dask.utils.Dispatch("dask_serialize")
 dask_deserialize = dask.utils.Dispatch("dask_deserialize")

distributed/protocol/tests/test_numpy.py

@@ -4,7 +4,7 @@
 
 np = pytest.importorskip("numpy")
 
-from dask.utils import tmpfile
+from dask.utils import ensure_bytes, tmpfile
 
 from distributed.protocol import (
     decompress,
@@ -20,7 +20,7 @@
 from distributed.protocol.pickle import HIGHEST_PROTOCOL
 from distributed.protocol.utils import BIG_BYTES_SHARD_SIZE
 from distributed.system import MEMORY_LIMIT
-from distributed.utils import ensure_bytes, nbytes
+from distributed.utils import nbytes
 from distributed.utils_test import gen_cluster
 
 

distributed/protocol/tests/test_pandas.py

@@ -4,6 +4,7 @@
 np = pytest.importorskip("numpy")
 
 from dask.dataframe.utils import assert_eq
+from dask.utils import ensure_bytes
 
 from distributed.protocol import (
     decompress,
@@ -13,7 +14,6 @@
     serialize,
     to_serialize,
 )
-from distributed.utils import ensure_bytes
 
 dfs = [
     pd.DataFrame({}),

distributed/tests/test_utils.py

@@ -1,4 +1,3 @@
-import array
 import asyncio
 import contextvars
 import functools
@@ -26,7 +25,6 @@
     LoopRunner,
     TimeoutError,
     _maybe_complex,
-    ensure_bytes,
     ensure_ip,
     format_dashboard_link,
     get_ip_interface,
@@ -248,27 +246,6 @@ def test_seek_delimiter_endline():
     assert f.tell() == 7
 
 
-def test_ensure_bytes():
-    data = [b"1", "1", memoryview(b"1"), bytearray(b"1"), array.array("b", [49])]
-    for d in data:
-        result = ensure_bytes(d)
-        assert isinstance(result, bytes)
-        assert result == b"1"
-
-
-def test_ensure_bytes_ndarray():
-    np = pytest.importorskip("numpy")
-    result = ensure_bytes(np.arange(12))
-    assert isinstance(result, bytes)
-
-
-def test_ensure_bytes_pyarrow_buffer():
-    pa = pytest.importorskip("pyarrow")
-    buf = pa.py_buffer(b"123")
-    result = ensure_bytes(buf)
-    assert isinstance(result, bytes)
-
-
 def test_nbytes():
     np = pytest.importorskip("numpy")
 

distributed/utils.py

@@ -47,6 +47,7 @@
 
 import dask
 from dask import istask
+from dask.utils import ensure_bytes as _ensure_bytes
 from dask.utils import parse_timedelta as _parse_timedelta
 from dask.widgets import get_template
 
@@ -1000,17 +1001,14 @@ def ensure_bytes(s):
     >>> ensure_bytes(b'123')
     b'123'
     """
-    if isinstance(s, bytes):
-        return s
-    elif hasattr(s, "encode"):
-        return s.encode()
-    else:
-        try:
-            return bytes(s)
-        except Exception as e:
-            raise TypeError(
-                "Object %s is neither a bytes object nor has an encode method" % s
-            ) from e
+    warnings.warn(
+        "`distributed.utils.ensure_bytes` is deprecated. "
+        "Please switch to `dask.utils.ensure_bytes`. "
+        "This will be removed in `2022.6.0`.",
+        DeprecationWarning,
+        stacklevel=2,
+    )
+    return _ensure_bytes(s)
 
 
 def open_port(host=""):