From e99033ee1a1aade648653f44295f0afd411092cb Mon Sep 17 00:00:00 2001 From: dcherian Date: Mon, 9 Mar 2020 15:05:05 +0530 Subject: [PATCH] Use list comprehension --- xarray/core/parallel.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/xarray/core/parallel.py b/xarray/core/parallel.py index 5d4a0661288..9263a0f809a 100644 --- a/xarray/core/parallel.py +++ b/xarray/core/parallel.py @@ -55,7 +55,7 @@ def check_result_variables( ) -def subset_dataset_to_chunk(graph, gname, dataset, input_chunks, chunk_tuple): +def subset_dataset_to_block(graph, gname, dataset, input_chunks, chunk_tuple): # mapping from dimension name to chunk index input_chunk_index = dict(zip(input_chunks.keys(), chunk_tuple)) @@ -420,16 +420,12 @@ def _wrapper(func, args, kwargs, arg_is_array, expected): # mapping from dimension name to chunk index input_chunk_index = dict(zip(input_chunks.keys(), chunk_tuple)) - chunked_args = [] - for arg in (dataset,) + tuple(converted_args): - if isinstance(arg, (DataArray, Dataset)): - chunked_args.append( - subset_dataset_to_chunk( - graph, gname, arg, input_chunks, chunk_tuple - ) - ) - else: - chunked_args.append(arg) + blocked_args = [ + subset_dataset_to_block(graph, gname, arg, input_chunks, chunk_tuple) + if isinstance(arg, (DataArray, Dataset)) + else arg + for arg in (dataset,) + tuple(converted_args) + ] # expected["shapes", "coords", "data_vars"] are used to raise nice error messages in _wrapper expected = {} @@ -447,7 +443,7 @@ def _wrapper(func, args, kwargs, arg_is_array, expected): graph[from_wrapper] = ( _wrapper, func, - chunked_args, + blocked_args, kwargs, [input_is_array] + arg_is_array, expected,