Function get_collection_names

dask/base.py:1331–1349

Infer the collection names from the dask keys, under the assumption that all keys are either tuples with matching first element, and that element is a string, or there is exactly one key and it is a string. Examples -------- >>> a.__dask_keys__() # doctest: +SKIP ["foo", "b

(collection)

def get_collection_names(collection) -> set[str]:
    """Infer the collection names from the dask keys, under the assumption that all keys
    are either tuples with matching first element, and that element is a string, or
    there is exactly one key and it is a string.

    Examples
    --------
    >>> a.__dask_keys__()  # doctest: +SKIP
    ["foo", "bar"]
    >>> get_collection_names(a)  # doctest: +SKIP
    {"foo", "bar"}
    >>> b.__dask_keys__()  # doctest: +SKIP
    [[("foo-123", 0, 0), ("foo-123", 0, 1)], [("foo-123", 1, 0), ("foo-123", 1, 1)]]
    >>> get_collection_names(b)  # doctest: +SKIP
    {"foo-123"}
    """
    if not is_dask_collection(collection):
        raise TypeError(f"Expected Dask collection; got {type(collection)}")
    return {get_name_from_key(k) for k in flatten(collection.__dask_keys__())}


def get_name_from_key(key: Key) -> str: