Use list comprehension

pydata · Mar 21, 2020 · e99033e · e99033e
1 parent 3035752
commit e99033e
Showing 1 changed file with 8 additions and 12 deletions.
diff --git 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,