From 24fb4d17dc1b77786596109954951c7117206abc Mon Sep 17 00:00:00 2001
From: LucaMarconato <2664412+LucaMarconato@users.noreply.github.com>
Date: Tue, 11 Feb 2025 16:27:28 +0100
Subject: [PATCH] fix suffix concatenate when each table annotates > 1 elements
 (#873)

fix suffix concatenate when concatenating 2 tables, each annotating 2 elements
---
 .github/workflows/test.yaml                   |  2 +-
 src/spatialdata/_core/concatenate.py          |  8 ++++-
 .../operations/test_spatialdata_operations.py | 32 +++++++++++++++++++
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml
index 1d7a79b9..788d0203 100644
--- a/.github/workflows/test.yaml
+++ b/.github/workflows/test.yaml
@@ -45,7 +45,7 @@ jobs:
               run: |
                   echo "::set-output name=dir::$(pip cache dir)"
             - name: Restore pip cache
-              uses: actions/cache@v2
+              uses: actions/cache@v4
               with:
                   path: ${{ steps.pip-cache-dir.outputs.dir }}
                   key: pip-${{ runner.os }}-${{ env.pythonLocation }}-${{ hashFiles('**/pyproject.toml') }}
diff --git a/src/spatialdata/_core/concatenate.py b/src/spatialdata/_core/concatenate.py
index 58e278f3..da1865d8 100644
--- a/src/spatialdata/_core/concatenate.py
+++ b/src/spatialdata/_core/concatenate.py
@@ -236,7 +236,13 @@ def _fix_ensure_unique_element_names(
             # fix the region_key column
             region, region_key, _ = get_table_keys(table)
             table.obs[region_key] = (table.obs[region_key].astype("str") + f"-{suffix}").astype("category")
-            table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY] = f"{region}-{suffix}"
+            new_region: str | list[str]
+            if isinstance(region, str):
+                new_region = f"{region}-{suffix}"
+            else:
+                assert isinstance(region, list)
+                new_region = [f"{r}-{suffix}" for r in region]
+            table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY] = new_region
 
             # fix the obs names
             if rename_obs_names:
diff --git a/tests/core/operations/test_spatialdata_operations.py b/tests/core/operations/test_spatialdata_operations.py
index 9b669fcb..dd765029 100644
--- a/tests/core/operations/test_spatialdata_operations.py
+++ b/tests/core/operations/test_spatialdata_operations.py
@@ -5,6 +5,7 @@
 import numpy as np
 import pytest
 from anndata import AnnData
+from geopandas import GeoDataFrame
 
 from spatialdata._core.concatenate import _concatenate_tables, concatenate
 from spatialdata._core.data_extent import are_extents_equal, get_extent
@@ -425,6 +426,37 @@ def test_concatenate_sdatas_from_iterable(concatenate_tables: bool, obs_names_ma
     assert sdata0["table"].obs_names[0] == "1"
 
 
+def test_concatenate_two_tables_each_annotating_two_elements() -> None:
+    # let's define 4 polygon elements. Two of them are annotated by the first table and the other two by the second
+    # table. The two tables have the same region key and instance key.
+    def _get_table_and_poly(i: int) -> tuple[AnnData, GeoDataFrame]:
+        poly = _get_shapes()["poly"]
+        n = len(poly)
+        region = f"poly{i}"
+        table = TableModel.parse(
+            AnnData(obs={"region": [region] * n, "instance_id": list(range(n))}),
+            region=region,
+            region_key="region",
+            instance_key="instance_id",
+        )
+        return table, poly
+
+    table0_a, poly0_a = _get_table_and_poly(0)
+    table1_a, poly1_a = _get_table_and_poly(1)
+    table0_b, poly0_b = _get_table_and_poly(0)
+    table1_b, poly1_b = _get_table_and_poly(1)
+
+    table01_a = _concatenate_tables([table0_a, table1_a])
+    table01_b = _concatenate_tables([table0_b, table1_b])
+
+    sdata01_a = SpatialData.init_from_elements({"poly0": poly0_a, "poly1": poly1_a, "table": table01_a})
+    sdata01_b = SpatialData.init_from_elements({"poly0": poly0_b, "poly1": poly1_b, "table": table01_b})
+    sdata = concatenate({"a": sdata01_a, "b": sdata01_b}, concatenate_tables=True)
+    region, _, _ = get_table_keys(sdata["table"])
+    assert region == ["poly0-a", "poly1-a", "poly0-b", "poly1-b"]
+    assert len(sdata["table"]) == 4 * len(poly0_a)
+
+
 def test_concatenate_sdatas_single_item() -> None:
     sdata = blobs()