Annotate "addresses" field as set

dask · Feb 10, 2021 · df6d8a4 · df6d8a4
1 parent e4c202a
commit df6d8a4
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
@@ -3524,10 +3524,12 @@ async def add_worker(
             if dh is None:
                 parent._host_info[host] = dh = dict()
 
-            if "addresses" not in dh:
-                dh.update({"addresses": set(), "nthreads": 0})
+            dh_addresses: set = dh.get("addresses")
+            if dh_addresses is None:
+                dh["addresses"] = dh_addresses = set()
+                dh["nthreads"] = 0
 
-            dh["addresses"].add(address)
+            dh_addresses.add(address)
             dh["nthreads"] += nthreads
 
             parent._total_nthreads += nthreads
@@ -4139,12 +4141,14 @@ async def remove_worker(self, comm=None, address=None, safe=False, close=True):
             if dh is None:
                 parent._host_info[host] = dh = dict()
 
+            dh_addresses: set = dh["addresses"]
+            dh_addresses.remove(address)
             dh["nthreads"] -= ws._nthreads
-            dh["addresses"].remove(address)
             parent._total_nthreads -= ws._nthreads
 
-            if not dh["addresses"]:
+            if not dh_addresses:
                 dh = None
+                dh_addresses = None
                 del parent._host_info[host]
 
             self.rpc.remove(address)