Repository: dask/dask
File: dask/array/rechunk.py (Python)

Code before:
    if method == "tasks":
        steps = plan_rechunk(
            x.chunks, chunks, x.dtype.itemsize, threshold, block_size_limit
        )
        for c in steps:
            x = _compute_rechunk(x, c)

        return x

    elif method == "p2p":
        from distributed.shuffle import rechunk_p2p

        return rechunk_p2p(
            x,
            chunks,
            threshold=threshold,
            block_size_limit=block_size_limit,
            balance=balance,
        )

    else:
        raise NotImplementedError(f"Unknown rechunking method '{method}'")


def _choose_rechunk_method(old_chunks, new_chunks, threshold=None):
    if method := config.get("array.rechunk.method", None) is not None:
        return method
    try:
        from distributed import default_client

        default_client()
    except (ImportError, ValueError):
        return "tasks"

    _old_to_new = old_to_new(old_chunks, new_chunks)
    graph_size = math.prod(sum(len(ins) for ins in axis) for axis in _old_to_new)
    threshold = threshold or config.get("array.rechunk.threshold")
    graph_size_threshold = _graph_size_threshold(old_chunks, new_chunks, threshold)
    return "tasks" if graph_size < graph_size_threshold else "p2p"


def _number_of_blocks(chunks):
    return reduce(mul, map(len, chunks))


def _largest_block_size(chunks):
    return reduce(mul, map(max, chunks))


def estimate_graph_size(old_chunks, new_chunks):
    """Estimate the graph size during a rechunk computation."""

Review comment (suggestion):
This is a little difficult to reason about. Maybe something like this instead?

```suggestion
    if config.get("array.rechunk.method", None) is not None:
```

Code after:
    if method == "tasks":
        steps = plan_rechunk(
            x.chunks, chunks, x.dtype.itemsize, threshold, block_size_limit
        )
        for c in steps:
            x = _compute_rechunk(x, c)

        return x

    elif method == "p2p":
        from distributed.shuffle import rechunk_p2p

        return rechunk_p2p(
            x,
            chunks,
            threshold=threshold,
            block_size_limit=block_size_limit,
            balance=balance,
        )

    else:
        raise NotImplementedError(f"Unknown rechunking method '{method}'")


def _choose_rechunk_method(old_chunks, new_chunks, threshold=None):
    if method := config.get("array.rechunk.method", None):
        return method
    try:
        from distributed import default_client

        default_client()
    except (ImportError, ValueError):
        return "tasks"

    _old_to_new = old_to_new(old_chunks, new_chunks)
    graph_size = math.prod(sum(len(ins) for ins in axis) for axis in _old_to_new)
    threshold = threshold or config.get("array.rechunk.threshold")
    graph_size_threshold = _graph_size_threshold(old_chunks, new_chunks, threshold)
    return "tasks" if graph_size < graph_size_threshold else "p2p"


def _number_of_blocks(chunks):
    return reduce(mul, map(len, chunks))


def _largest_block_size(chunks):
    return reduce(mul, map(max, chunks))


def estimate_graph_size(old_chunks, new_chunks):
    """Estimate the graph size during a rechunk computation."""