Repository: dask/dask
File: dask/_expr.py (Python)

Code before:
        if isinstance(key, list):
            new_keys.append(_convert_dask_keys(key))
        else:
            new_keys.append(TaskRef(key))
    return List(*new_keys)


class HLGFinalizeCompute(HLGExpr):

    def _simplify_down(self):
        if not self.postcompute:
            return self.dsk

        from dask.delayed import Delayed

        # Skip finalization for Delayed
        if self.dsk.postcompute == Delayed.__dask_postcompute__(self.dsk):
            return self.dsk
        return self

    @property
    def _name(self):
        return f"finalize-{self.deterministic_token}"

    def __dask_graph__(self):
        # The baseclass __dask_graph__ will not just materialie this layer but
        # also that of its dependencies, i.e. it will render the finalized and
        # the non-finalized graph and combine them. We only want the finalized
        # so we're overriding this.
        # This is an artifact generated since the wrapped expression is
        # identified automatically as a dependency but HLG expressions are not
        # working in this layered way.
        return self._layer()

    @property
    def hlg(self):
        expr = self.operand("dsk")
        layers = expr.dsk.layers.copy()
        deps = expr.dsk.dependencies.copy()
        keys = expr.__dask_keys__()
        func, extra_args = expr.postcompute
        t = Task(self._name, func, _convert_dask_keys(keys), *extra_args)
        from dask.highlevelgraph import HighLevelGraph, MaterializedLayer

        layers[t.key] = MaterializedLayer({t.key: t})
        deps[t.key] = set(expr.dsk.dependencies)
        return HighLevelGraph(layers, dependencies=deps)

    def __dask_keys__(self):
        return [self._name]


Review comment (suggestion):
nit
```suggestion
        # The baseclass __dask_graph__ will not just materialize this layer but
```

Code after:
            new_keys.append(_convert_dask_keys(key))
        else:
            new_keys.append(TaskRef(key))
    return List(*new_keys)


class HLGFinalizeCompute(HLGExpr):

    def _simplify_down(self):
        if not self.postcompute:
            return self.dsk

        from dask.delayed import Delayed

        # Skip finalization for Delayed
        if self.dsk.postcompute == Delayed.__dask_postcompute__(self.dsk):
            return self.dsk
        return self

    @property
    def _name(self):
        return f"finalize-{self.deterministic_token}"

    def __dask_graph__(self):
        # The baseclass __dask_graph__ will not just materialize this layer but
        # also that of its dependencies, i.e. it will render the finalized and
        # the non-finalized graph and combine them. We only want the finalized
        # so we're overriding this.
        # This is an artifact generated since the wrapped expression is
        # identified automatically as a dependency but HLG expressions are not
        # working in this layered way.
        return self._layer()

    @property
    def hlg(self):
        expr = self.operand("dsk")
        layers = expr.dsk.layers.copy()
        deps = expr.dsk.dependencies.copy()
        keys = expr.__dask_keys__()
        func, extra_args = expr.postcompute
        t = Task(self._name, func, _convert_dask_keys(keys), *extra_args)
        from dask.highlevelgraph import HighLevelGraph, MaterializedLayer

        layers[t.key] = MaterializedLayer({t.key: t})
        deps[t.key] = set(expr.dsk.dependencies)
        return HighLevelGraph(layers, dependencies=deps)

    def __dask_keys__(self):
        return [self._name]