Function cull

dask/_task_spec.py:1188–1210

(
    dsk: dict[KeyType, GraphNode], keys: Iterable[KeyType]
)

def cull(
    dsk: dict[KeyType, GraphNode], keys: Iterable[KeyType]
) -> dict[KeyType, GraphNode]:
    if not isinstance(keys, (list, set, tuple)):
        raise TypeError(
            f"Expected list, set or tuple for keys, got {type(keys).__name__}"
        )
    if len(keys) == len(dsk):
        return dsk
    work = set(keys)
    seen: set[KeyType] = set()
    dsk2 = {}
    wpop = work.pop
    wupdate = work.update
    sadd = seen.add
    while work:
        k = wpop()
        if k in seen or k not in dsk:
            continue
        sadd(k)
        dsk2[k] = v = dsk[k]
        wupdate(v.dependencies)
    return dsk2


@functools.cache