Implement name=

crusaderky · crusaderky · commit 16fbb10ae8cc · 2021-01-25T10:35:21.000Z
diff --git a/dask/bag/core.py b/dask/bag/core.py
@@ -40,7 +40,14 @@
 from ..base import tokenize, dont_optimize, DaskMethodsMixin
 from ..bytes import open_files
 from ..context import globalmethod
-from ..core import quote, istask, get_dependencies, reverse_dict, flatten
+from ..core import (
+    quote,
+    istask,
+    get_dependencies,
+    replace_name_in_key,
+    reverse_dict,
+    flatten,
+)
 from ..sizeof import sizeof
 from ..delayed import Delayed, unpack_collections
 from ..highlevelgraph import HighLevelGraph
@@ -365,7 +372,9 @@ def __dask_postpersist__(self):
         return self._rebuild, ()
 
     def _rebuild(self, dsk, name=None):
-        key = change_name_in_key(self.key, name) if name else self.key
+        key = self.key
+        if name:
+            key = replace_name_in_key(key, name)
         return Item(dsk, key)
 
     @staticmethod
diff --git a/dask/core.py b/dask/core.py
@@ -503,3 +503,21 @@ def quote(x):
     if istask(x) or type(x) is list or type(x) is dict:
         return (literal(x),)
     return x
+
+
+def replace_name_in_key(key, name: str):
+    """Given a dask key, which must be either a single string or a tuple whose first
+    element is a string (commonly referred to as 'name'), replace the name with a new
+    one.
+
+    e.g.::
+
+      >>> replace_name_in_key("foo", "bar")
+      "bar"
+      >>> replace_name_in_key(("foo", 1, 2), "bar") -> ("bar", 1, 2)
+    """
+    if isinstance(key, str):
+        return name
+    if isinstance(key, tuple) and key and isinstance(key[0], str):
+        return (name,) + key[1:]
+    raise TypeError(f"Expected str or tuple[str, Hashable, ...]; got {key}")
diff --git a/dask/delayed.py b/dask/delayed.py
@@ -11,7 +11,7 @@
 from .base import tokenize as _tokenize
 from .compatibility import is_dataclass, dataclass_fields
 
-from .core import quote
+from .core import replace_name_in_key, quote
 from .context import globalmethod
 from .optimization import cull
 from .utils import funcname, methodcaller, OperatorMethodMixin, ensure_dict, apply
@@ -512,7 +512,9 @@ def __dask_postpersist__(self):
         return self._rebuild, ()
 
     def _rebuild(self, dsk, name=None):
-        key = change_name_in_key(self.key, name) if name else self.key
+        key = self.key
+        if name:
+            key = replace_name_in_key(key, name)
         return Delayed(key, dsk, getattr(self, "_length", None))
 
     def __getstate__(self):
diff --git a/dask/graphsurgery.py b/dask/graphsurgery.py
@@ -5,11 +5,12 @@
 from .array import Array
 from .bag import Bag
 from .base import tokenize, unpack_collections
+from .blockwise import blockwise
 from .core import flatten
 from .dataframe import DataFrame, Series
 from .delayed import Delayed, delayed
-from .highlevelgraph import BasicLayer, Layer, HighLevelGraph
-from .blockwise import blockwise
+from .highlevelgraph import BasicLayer, HighLevelGraph, Layer
+from .utils import ensure_dict
 
 __all__ = ("bind", "choke", "clone", "drop")
 
@@ -206,7 +207,7 @@ def _bind_one_by_layers(
 def _bind_one_by_keys(
     child: T, blocker: Optional[Delayed], omit_keys: set, seed: Hashable
 ) -> T:
-    dsk = dict(child.__dask_graph__())
+    dsk = ensure_dict(child.__dask_graph__())
 
     # TODO
 
