Import tlz (#3579)

Import from `tlz` for optional `cytoolz` support

distributed/cfexecutor.py

@@ -1,7 +1,7 @@
 import concurrent.futures as cf
 import weakref
 
-from toolz import merge
+from tlz import merge
 
 from tornado import gen
 

distributed/cli/dask_worker.py

@@ -22,7 +22,7 @@
 )
 from distributed.utils import deserialize_for_cli, import_term
 
-from toolz import valmap
+from tlz import valmap
 from tornado.ioloop import IOLoop, TimeoutError
 
 logger = logging.getLogger("distributed.dask_worker")

distributed/client.py

@@ -30,10 +30,8 @@
 from dask.compatibility import apply
 from dask.utils import ensure_dict, format_bytes, funcname
 
-try:
-    from cytoolz import first, groupby, merge, valmap, keymap
-except ImportError:
-    from toolz import first, groupby, merge, valmap, keymap
+from tlz import first, groupby, merge, valmap, keymap
+
 try:
     from dask.delayed import single_key
 except ImportError:

distributed/core.py

@@ -10,7 +10,7 @@
 
 import dask
 import tblib
-from toolz import merge
+from tlz import merge
 from tornado import gen
 from tornado.ioloop import IOLoop
 

distributed/dashboard/components/__init__.py

@@ -26,7 +26,6 @@
 from bokeh.plotting import figure
 import dask
 from tornado import gen
-import toolz
 
 from distributed.dashboard.utils import without_property_validation, BOKEH_VERSION
 from distributed import profile

distributed/dashboard/components/scheduler.py

@@ -35,7 +35,8 @@
 from bokeh.io import curdoc
 import dask
 from dask.utils import format_bytes, key_split
-from toolz import pipe
+from tlz import pipe
+from tlz.curried import map, concat, groupby
 from tornado import escape
 
 try:
@@ -63,11 +64,6 @@
 from distributed.diagnostics.graph_layout import GraphLayout
 from distributed.diagnostics.task_stream import TaskStreamPlugin
 
-try:
-    from cytoolz.curried import map, concat, groupby
-except ImportError:
-    from toolz.curried import map, concat, groupby
-
 if dask.config.get("distributed.dashboard.export-tool"):
     from distributed.dashboard.export_tool import ExportTool
 else:

distributed/dashboard/components/shared.py

@@ -15,7 +15,7 @@
 from bokeh.plotting import figure
 import dask
 from tornado import gen
-import toolz
+import tlz as toolz
 
 from distributed.dashboard.components import DashboardComponent
 from distributed.dashboard.utils import (

distributed/dashboard/components/worker.py

@@ -20,7 +20,7 @@
 from bokeh.palettes import RdBu
 from bokeh.themes import Theme
 from dask.utils import format_bytes
-from toolz import merge, partition_all
+from tlz import merge, partition_all
 
 from distributed.dashboard.components import add_periodic_callback
 from distributed.dashboard.components.shared import (

distributed/dashboard/scheduler.py

@@ -8,10 +8,7 @@
 import dask
 from dask.utils import format_bytes
 
-try:
-    from cytoolz import merge, merge_with
-except ImportError:
-    from toolz import merge, merge_with
+from tlz import merge, merge_with
 
 from tornado import escape
 from tornado.websocket import WebSocketHandler

distributed/dashboard/tests/test_scheduler_bokeh.py

@@ -7,7 +7,7 @@
 import pytest
 
 pytest.importorskip("bokeh")
-from toolz import first
+from tlz import first
 from tornado import gen
 from tornado.httpclient import AsyncHTTPClient, HTTPRequest
 
@@ -624,7 +624,7 @@ def test_proxy_to_workers(c, s, a, b):
     },
 )
 async def test_lots_of_tasks(c, s, a, b):
-    import toolz
+    import tlz as toolz
 
     ts = TaskStream(s)
     ts.update()

distributed/dashboard/tests/test_worker_bokeh.py

@@ -6,7 +6,7 @@
 
 pytest.importorskip("bokeh")
 import sys
-from toolz import first
+from tlz import first
 from tornado import gen
 from tornado.httpclient import AsyncHTTPClient
 

distributed/dashboard/utils.py

@@ -5,19 +5,15 @@
 import bokeh
 from bokeh.io import curdoc
 from tornado import web
-from toolz import partition
+from tlz import partition
+from tlz.curried import first
 
 try:
     import numpy as np
 except ImportError:
     np = False
 
 
-try:
-    from cytoolz.curried import first
-except ImportError:
-    from toolz.curried import first
-
 BOKEH_VERSION = LooseVersion(bokeh.__version__)
 dirname = os.path.dirname(__file__)
 

distributed/dashboard/worker.py

@@ -3,7 +3,7 @@
 import os
 
 from bokeh.themes import Theme
-from toolz import merge
+from tlz import merge
 
 from .components.worker import (
     status_doc,

distributed/deploy/adaptive_core.py

@@ -2,7 +2,7 @@
 import math
 
 from tornado.ioloop import IOLoop
-import toolz
+import tlz as toolz
 
 from ..metrics import time
 from ..utils import parse_timedelta, PeriodicCallback

distributed/deploy/old_ssh.py

@@ -12,7 +12,7 @@
 
 from threading import Thread
 
-from toolz import merge
+from tlz import merge
 
 from tornado import gen
 

distributed/deploy/tests/test_spec_cluster.py

@@ -8,7 +8,7 @@
 from distributed.metrics import time
 from distributed.utils_test import loop, cleanup  # noqa: F401
 from distributed.utils import is_valid_xml
-import toolz
+import tlz as toolz
 import pytest
 
 

distributed/diagnostics/progress.py

@@ -3,7 +3,7 @@
 import logging
 from timeit import default_timer
 
-from toolz import groupby, valmap
+from tlz import groupby, valmap
 
 from .plugin import SchedulerPlugin
 from ..utils import key_split, key_split_group, log_errors, tokey

distributed/diagnostics/progress_stream.py

@@ -1,6 +1,6 @@
 import logging
 
-from toolz import valmap, merge
+from tlz import valmap, merge
 
 from .progress import AllProgress
 

distributed/diagnostics/progressbar.py

@@ -4,7 +4,7 @@
 import sys
 import weakref
 
-from toolz import valmap
+from tlz import valmap
 from tornado.ioloop import IOLoop
 
 from .progress import format_time, Progress, MultiProgress

distributed/diagnostics/tests/test_task_stream.py

@@ -2,7 +2,7 @@
 from time import sleep
 
 import pytest
-from toolz import frequencies
+from tlz import frequencies
 
 from distributed import get_task_stream
 from distributed.utils_test import gen_cluster, div, inc, slowinc

distributed/diagnostics/tests/test_widgets.py

@@ -74,7 +74,7 @@ def record_display(*args):
 from operator import add
 import re
 
-from toolz import valmap
+from tlz import valmap
 
 from distributed.client import wait
 from distributed.worker import dumps_task

distributed/profile.py

@@ -31,7 +31,7 @@
 import threading
 from time import sleep
 
-import toolz
+import tlz as toolz
 
 from .metrics import time
 from .utils import format_time, color_of, parse_timedelta

distributed/protocol/compression.py

@@ -3,11 +3,12 @@
 
 Includes utilities for determining whether or not to compress
 """
+from functools import partial
 import logging
 import random
 
 import dask
-from toolz import identity, partial
+from tlz import identity
 
 try:
     import blosc

distributed/protocol/core.py

@@ -1,13 +1,9 @@
+from functools import reduce
 import logging
 import operator
 
 import msgpack
 
-try:
-    from cytoolz import reduce
-except ImportError:
-    from toolz import reduce
-
 from .compression import compressions, maybe_compress, decompress
 from .serialize import serialize, deserialize, Serialize, Serialized, extract_serialize
 from .utils import frame_split_size, merge_frames, msgpack_opts

distributed/protocol/serialize.py

@@ -4,10 +4,7 @@
 import dask
 from dask.base import normalize_token
 
-try:
-    from cytoolz import valmap, get_in
-except ImportError:
-    from toolz import valmap, get_in
+from tlz import valmap, get_in
 
 import msgpack
 

distributed/protocol/tests/test_serialize.py

@@ -4,7 +4,7 @@
 import msgpack
 import numpy as np
 import pytest
-from toolz import identity
+from tlz import identity
 
 from distributed import wait
 from distributed.protocol import (

distributed/scheduler.py

@@ -19,11 +19,18 @@
 import psutil
 import sortedcontainers
 
-try:
-    from cytoolz import frequencies, merge, pluck, merge_sorted, first, merge_with
-except ImportError:
-    from toolz import frequencies, merge, pluck, merge_sorted, first, merge_with
-from toolz import valmap, second, compose, groupby
+from tlz import (
+    frequencies,
+    merge,
+    pluck,
+    merge_sorted,
+    first,
+    merge_with,
+    valmap,
+    second,
+    compose,
+    groupby,
+)
 from tornado.ioloop import IOLoop
 
 import dask

distributed/stealing.py

@@ -9,10 +9,7 @@
 from .diagnostics.plugin import SchedulerPlugin
 from .utils import log_errors, parse_timedelta, PeriodicCallback
 
-try:
-    from cytoolz import topk
-except ImportError:
-    from toolz import topk
+from tlz import topk
 
 LATENCY = 10e-3
 log_2 = log(2)

distributed/tests/test_batched.py

@@ -2,7 +2,7 @@
 import random
 
 import pytest
-from toolz import assoc
+from tlz import assoc
 
 from distributed.batched import BatchedSend
 from distributed.core import listen, connect, CommClosedError

distributed/tests/test_client.py

@@ -1,5 +1,6 @@
 import asyncio
 from collections import deque
+from functools import partial
 import gc
 import logging
 from operator import add
@@ -18,7 +19,7 @@
 import zipfile
 
 import pytest
-from toolz import identity, isdistinct, concat, pluck, valmap, partial, first, merge
+from tlz import identity, isdistinct, concat, pluck, valmap, first, merge
 from tornado import gen
 
 import dask

distributed/tests/test_client_executor.py

@@ -11,7 +11,7 @@
 )
 
 import pytest
-from toolz import take
+from tlz import take
 
 from distributed import Client
 from distributed.utils import CancelledError

distributed/tests/test_failed_workers.py

@@ -3,7 +3,7 @@
 from time import sleep
 
 import pytest
-from toolz import partition_all, first
+from tlz import partition_all, first
 from tornado import gen
 
 from dask import delayed

distributed/tests/test_ipython.py

@@ -1,7 +1,7 @@
 from unittest import mock
 
 import pytest
-from toolz import first
+from tlz import first
 import tornado
 
 from distributed import Client

distributed/tests/test_nanny.py

@@ -9,7 +9,7 @@
 import numpy as np
 
 import pytest
-from toolz import valmap, first
+from tlz import valmap, first
 from tornado import gen
 from tornado.ioloop import IOLoop
 

distributed/tests/test_profile.py

@@ -1,7 +1,7 @@
 import pytest
 import sys
 import time
-from toolz import first
+from tlz import first
 import threading
 
 from distributed.compatibility import WINDOWS