diff --git a/distributed/comm/ws.py b/distributed/comm/ws.py index 2f6a9279fd8..dcfdc41e8d0 100644 --- a/distributed/comm/ws.py +++ b/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__) diff --git a/distributed/protocol/tests/test_numpy.py b/distributed/protocol/tests/test_numpy.py index 93c42aacf1d..895d5ad8896 100644 --- a/distributed/protocol/tests/test_numpy.py +++ b/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 diff --git a/distributed/protocol/tests/test_pandas.py b/distributed/protocol/tests/test_pandas.py index 58bfb90f75e..07fc916064d 100644 --- a/distributed/protocol/tests/test_pandas.py +++ b/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({}), diff --git a/distributed/tests/test_utils.py b/distributed/tests/test_utils.py index f50bd2c0081..ed1d275907a 100644 --- a/distributed/tests/test_utils.py +++ b/distributed/tests/test_utils.py @@ -26,7 +26,6 @@ LoopRunner, TimeoutError, _maybe_complex, - ensure_bytes, ensure_ip, ensure_memoryview, format_dashboard_link, @@ -249,27 +248,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_ensure_memoryview_empty(): result = ensure_memoryview(b"") assert isinstance(result, memoryview) diff --git a/distributed/utils.py b/distributed/utils.py index 475ea836aa5..4e7fe43dc43 100644 --- a/distributed/utils.py +++ b/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 ensure_memoryview(obj):