caught empty list of partitions in xu.merge_partitions()

tests/test_partitioning.py

@@ -132,10 +132,15 @@ def test_partition_roundtrip(self):
         reordered = back.isel(mesh2d_nFaces=order)
         assert reordered["face_z"].equals(self.uds["face_z"])
 
+    def test_merge_partition_single(self):
+        partitions = [self.uds]
+        back = pt.merge_partitions(partitions)
+        assert back == self.uds
+
     def test_merge_partitions__errors(self):
         partitions = self.uds.ugrid.partition(n_part=2)
         with pytest.raises(TypeError, match="Expected UgridDataArray or UgridDataset"):
-            pt.merge_partitions(p.ugrid.obj for p in partitions)
+            pt.merge_partitions([p.ugrid.obj for p in partitions])
 
         grid1 = partitions[1].ugrid.grid
         partitions[1]["extra"] = (grid1.face_dimension, np.ones(grid1.n_face))
@@ -162,6 +167,11 @@ def test_merge_partitions__errors(self):
         ):
             pt.merge_partitions(partitions)
 
+        with pytest.raises(
+            ValueError, match="Cannot merge partitions: zero partitions provided."
+        ):
+            xu.merge_partitions([])
+
     def test_merge_partitions_no_duplicates(self):
         part1 = self.uds.isel(mesh2d_nFaces=[0, 1, 2, 3])
         part2 = self.uds.isel(mesh2d_nFaces=[2, 3, 4, 5])

xugrid/ugrid/partitioning.py

@@ -328,6 +328,8 @@ def merge_partitions(partitions, merge_ugrid_chunks: bool = True):
     -------
     merged : UgridDataset
     """
+    if len(partitions) == 0:
+        raise ValueError("Cannot merge partitions: zero partitions provided.")
     types = {type(obj) for obj in partitions}
     msg = "Expected UgridDataArray or UgridDataset, received: {}"
     if len(types) > 1:
@@ -337,6 +339,10 @@ def merge_partitions(partitions, merge_ugrid_chunks: bool = True):
     if obj_type not in (UgridDataArray, UgridDataset):
         raise TypeError(msg.format(obj_type.__name__))
 
+    # return first partition if single partition is provided
+    if len(partitions) == 1:
+        return next(iter(partitions))
+
     # Collect grids
     grids = [grid for p in partitions for grid in p.grids]
     ugrid_dims = {dim for grid in grids for dim in grid.dimensions}