fixed bug with chunks passed as dict

Author: TomNicholas

xarray/namedarray/core.py

@@ -812,6 +812,8 @@ def chunk(
             chunks = either_dict_or_kwargs(chunks, chunks_kwargs, "chunk")
 
         if is_dict_like(chunks):
+            # turns dict[str, tuple[in, ..]] -> dict[int, tuple[int, ...]]
+
             # This method of iteration allows for duplicated dimension names, GH8579
             chunks = {
                 dim_number: chunks[dim]
@@ -825,6 +827,21 @@ def chunk(
             # if is_dict_like(chunks) and chunks != {}:
             #     chunks = tuple(chunks.get(n, s) for n, s in enumerate(data_old.shape))  # type: ignore[assignment]
 
+            old_chunks = data_old.chunks
+
+            if is_dict_like(chunks):
+                # turns dict[int, tuple[int, ...]] -> tuple[tuple[int, ...], ...], filling in unspecified dimensions using previous chunking
+                chunks = tuple(
+                    [
+                        (
+                            chunks[dim_number]
+                            if dim_number in chunks.keys()
+                            else old_chunks[dim_number]
+                        )
+                        for dim_number in range(self.ndim)
+                    ]
+                )
+
             print(f"hopefully normalized chunks = {chunks}")
 
             # Assume any chunked array supports .rechunk - if it doesn't then at least a clear AttributeError will be raised.
@@ -846,6 +863,7 @@ def chunk(
                 ndata = ImplicitToExplicitIndexingAdapter(data_old, OuterIndexer)  # type: ignore[assignment]
 
             if is_dict_like(chunks):
+                # turns dict[int, tuple[int, ...]] -> tuple[tuple[int, ...], ...], filling in unspecified dimensions using full length along each axis (i.e. array shape)
                 chunks = tuple(chunks.get(n, s) for n, s in enumerate(ndata.shape))  # type: ignore[assignment]
 
             chunkmanager = guess_chunkmanager(chunked_array_type)