From 24b61820980350ee5bdc10bdfa3b2035219ff369 Mon Sep 17 00:00:00 2001
From: veenstrajelmer <60435591+veenstrajelmer@users.noreply.github.com>
Date: Wed, 14 Aug 2024 13:47:18 +0200
Subject: [PATCH] escape for single partition

---
 tests/test_partitioning.py   | 5 +++++
 xugrid/ugrid/partitioning.py | 5 +++++
 2 files changed, 10 insertions(+)

diff --git a/tests/test_partitioning.py b/tests/test_partitioning.py
index 4a3a7654..274bb87c 100644
--- a/tests/test_partitioning.py
+++ b/tests/test_partitioning.py
@@ -132,6 +132,11 @@ 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 = xu.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"):
diff --git a/xugrid/ugrid/partitioning.py b/xugrid/ugrid/partitioning.py
index 519cf74d..e68d30a9 100644
--- a/xugrid/ugrid/partitioning.py
+++ b/xugrid/ugrid/partitioning.py
@@ -305,6 +305,7 @@ def merge_data_along_dim(
 
     return xr.concat(to_merge, dim=merge_dim)
 
+xu.merge_partitions([1])
 
 def merge_partitions(partitions, merge_ugrid_chunks: bool = True):
     """
@@ -339,6 +340,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}