Remove sortedcontainers entirely

gjoseph92 · gjoseph92 · commit aa0cde06c0a4 · 2022-11-02T21:34:27.000-06:00
`Computation.code` was the only other place it was used. Doesn't seem worth the dependency for that.
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
@@ -34,7 +34,6 @@
 from typing import TYPE_CHECKING, Any, ClassVar, Literal, cast, overload
 
 import psutil
-from sortedcontainers import SortedSet
 from tlz import (
     first,
     groupby,
@@ -854,15 +853,15 @@ class Computation:
 
     start: float
     groups: set[TaskGroup]
-    code: SortedSet
+    code: dict[str, None]  # a sorted set
     id: uuid.UUID
 
     __slots__ = tuple(__annotations__)
 
     def __init__(self):
         self.start = time()
         self.groups = set()
-        self.code = SortedSet()
+        self.code = {}
         self.id = uuid.uuid4()
 
     @property
@@ -4268,8 +4267,8 @@ def update_graph(
             computation = Computation()
             self.computations.append(computation)
 
-        if code and code not in computation.code:  # add new code blocks
-            computation.code.add(code)
+        if code:  # add new code blocks
+            computation.code.setdefault(code, None)
 
         n = 0
         while len(tasks) != n:  # walk through new tasks, cancel any bad deps
diff --git a/distributed/tests/test_client.py b/distributed/tests/test_client.py
@@ -6985,7 +6985,7 @@ def fetch_comp_code(dask_scheduler):
         assert len(computations) == 1
         comp = computations[0]
         assert len(comp.code) == 1
-        return comp.code[0]
+        return first(comp.code)
 
     code = client.run_on_scheduler(fetch_comp_code)
 
@@ -7005,7 +7005,7 @@ def fetch_comp_code(dask_scheduler):
         assert len(computations) == 1
         comp = computations[0]
         assert len(comp.code) == 1
-        return comp.code[0]
+        return first(comp.code)
 
     code = client.run_on_scheduler(fetch_comp_code)
     assert code == "<Code not available>"
@@ -7026,7 +7026,7 @@ async def test_computation_object_code_dask_persist(c, s, a, b):
     comp = computations[0]
     assert len(comp.code) == 1
 
-    assert comp.code[0] == test_function_code
+    assert first(comp.code) == test_function_code
 
 
 @gen_cluster(client=True)
@@ -7047,7 +7047,7 @@ def func(x):
 
     assert len(comp.code) == 1
 
-    assert comp.code[0] == test_function_code
+    assert first(comp.code) == test_function_code
 
 
 @gen_cluster(client=True)
@@ -7069,7 +7069,7 @@ def func(x):
     # Code is deduplicated
     assert len(comp.code) == 1
 
-    assert comp.code[0] == test_function_code
+    assert first(comp.code) == test_function_code
 
 
 @gen_cluster(client=True)
@@ -7091,7 +7091,7 @@ def func(x):
     # Code is deduplicated
     assert len(comp.code) == 1
 
-    assert comp.code[0] == test_function_code
+    assert first(comp.code) == test_function_code
 
 
 @gen_cluster(client=True)
@@ -7109,7 +7109,7 @@ async def test_computation_object_code_client_map(c, s, a, b):
     comp = computations[0]
     assert len(comp.code) == 1
 
-    assert comp.code[0] == test_function_code
+    assert first(comp.code) == test_function_code
 
 
 @gen_cluster(client=True)
@@ -7127,7 +7127,7 @@ async def test_computation_object_code_client_compute(c, s, a, b):
     comp = computations[0]
     assert len(comp.code) == 1
 
-    assert comp.code[0] == test_function_code
+    assert first(comp.code) == test_function_code
 
 
 @gen_cluster(client=True, Worker=Nanny)
diff --git a/requirements.txt b/requirements.txt
@@ -6,7 +6,6 @@ locket >= 1.0.0
 msgpack >= 0.6.0
 packaging >= 20.0
 psutil >= 5.0
-sortedcontainers !=2.0.0, !=2.0.1
 tblib >= 1.6.0
 toolz >= 0.8.2
 tornado >= 6.0.3, <6.2
