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

Code before:
                io_func = subgraph.get(any_key)[0]
            else:
                io_func = None
            ninds = 1 if isinstance(output_indices, str) else len(output_indices)
            dsk = {
                output: (
                    PackedFunctionCall(io_func),
                    *[blockwise_token(i) for i in range(ninds)],
                )
            }

        # Super-class initializer
        super().__init__(
            output,
            output_indices,
            dsk,
            indices,
            numblocks,
            concatenate=None,
            new_axes=None,
            output_blocks=None,
        )

        # BlockwiseIO requires an `io_subgraph` input.
        # This is a tuple: (key-name, subgraph)
        self.io_name, self.io_subgraph = io_subgraph

    @property
    def _dict(self):
        if hasattr(self, "_cached_dict"):
            return self._cached_dict["dsk"]
        else:
            keys = tuple(map(blockwise_token, range(len(self.indices))))
            dsk, _ = fuse(self.dsk, [self.output])
            func = SubgraphCallable(dsk, self.output, keys)

            dsk = make_blockwise_graph(
                func,
                self.output,
                self.output_indices,
                *list(toolz.concat(self.indices)),
                new_axes=self.new_axes,
                numblocks=self.numblocks,
                concatenate=self.concatenate,
                output_blocks=self.output_blocks,
                dims=self.dims,
            )

            # Handle IO Subgraph
            for k in dsk:
                io_key = (self.io_name,) + tuple([k[i] for i in range(1, len(k))])

Review comment (suggestion):
Not sure if it makes sense to continue accepting a `(name, io_subgraph)` tuple in the `BlockwiseIO` initializer.  I personally like this, but also recognize that two separate args may be more intuitive.

Code after:
        if not dsk:
            # Extract actual IO function for SubgraphCallable construction.
            # Wrap func in `PackedFunctionCall`, since it will receive
            # all arguments as a sigle (packed) tuple at run time.
            if io_subgraph:
                # We assume a 1-to-1 mapping between keys (i.e. tasks) and
                # chunks/partitions in `io_subgraph`, and assume the first
                # (callable) element is the same for all tasks.
                any_key = next(iter(io_subgraph))
                io_func = io_subgraph.get(any_key)[0]
            else:
                io_func = None
            ninds = 1 if isinstance(output_indices, str) else len(output_indices)
            dsk = {
                output: (
                    PackedFunctionCall(io_func),
                    *[blockwise_token(i) for i in range(ninds)],
                )
            }

        # Super-class initializer
        super().__init__(
            output,
            output_indices,
            dsk,
            indices,
            numblocks,
            concatenate=None,
            new_axes=None,
            output_blocks=None,
            annotations=annotations,
        )

        # BlockwiseIO requires `io_name` and `io_subgraph` inputs
        self.io_name = io_name
        self.io_subgraph = io_subgraph

    @property
    def _dict(self):
        if hasattr(self, "_cached_dict"):
            return self._cached_dict["dsk"]
        else:
            keys = tuple(map(blockwise_token, range(len(self.indices))))
            dsk, _ = fuse(self.dsk, [self.output])
            func = SubgraphCallable(dsk, self.output, keys)

            dsk = make_blockwise_graph(
                func,
                self.output,
                self.output_indices,
                *list(toolz.concat(self.indices)),