@@ -55,7 +55,7 @@ def check_result_variables(
55
55
)
56
56
57
57
58
- def subset_dataset_to_chunk (graph , gname , dataset , input_chunks , chunk_tuple ):
58
+ def subset_dataset_to_block (graph , gname , dataset , input_chunks , chunk_tuple ):
59
59
60
60
# mapping from dimension name to chunk index
61
61
input_chunk_index = dict (zip (input_chunks .keys (), chunk_tuple ))
@@ -420,16 +420,12 @@ def _wrapper(func, args, kwargs, arg_is_array, expected):
420
420
# mapping from dimension name to chunk index
421
421
input_chunk_index = dict (zip (input_chunks .keys (), chunk_tuple ))
422
422
423
- chunked_args = []
424
- for arg in (dataset ,) + tuple (converted_args ):
425
- if isinstance (arg , (DataArray , Dataset )):
426
- chunked_args .append (
427
- subset_dataset_to_chunk (
428
- graph , gname , arg , input_chunks , chunk_tuple
429
- )
430
- )
431
- else :
432
- chunked_args .append (arg )
423
+ blocked_args = [
424
+ subset_dataset_to_block (graph , gname , arg , input_chunks , chunk_tuple )
425
+ if isinstance (arg , (DataArray , Dataset ))
426
+ else arg
427
+ for arg in (dataset ,) + tuple (converted_args )
428
+ ]
433
429
434
430
# expected["shapes", "coords", "data_vars"] are used to raise nice error messages in _wrapper
435
431
expected = {}
@@ -447,7 +443,7 @@ def _wrapper(func, args, kwargs, arg_is_array, expected):
447
443
graph [from_wrapper ] = (
448
444
_wrapper ,
449
445
func ,
450
- chunked_args ,
446
+ blocked_args ,
451
447
kwargs ,
452
448
[input_is_array ] + arg_is_array ,
453
449
expected ,
0 commit comments