Repository: scverse/spatialdata

Issue:
`set_table_annotates_spatialelement` doesn't return a valid table
https://github.com/scverse/spatialdata/blob/9e07583f13a70be0855c94ce54b18f133f5ab5b8/src/spatialdata/_core/spatialdata.py#L432

using this, it returns a table that is not valid. Specifically, it fails at write time due to region not being categorical. I believe because the table is never actually parsed somehow, but important to check where does this check fail

cc. @lotte-vandevreken @lopollar

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/scverse/spatialdata/pull/610?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 91.94%. Comparing base [(`0727b08`)](https://app.codecov.io/gh/scverse/spatialdata/commit/0727b0810c5ac09c5b9ee2241b32302e00f5c1ac?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) to head [(`710ec03`)](https://app.codecov.io/gh/scverse/spatialdata/commit/710ec03f715a7ae1cadb387739f77d589b63ef42?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse).
> Report is 4 commits behind head on main.

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main     #610      +/-   ##
==========================================
+ Coverage   91.91%   91.94%   +0.02%     
==========================================
  Files          44       44              
  Lines        6641     6651      +10     
==========================================
+ Hits         6104     6115      +11     
+ Misses        537      536       -1     
```

| [Files](https://app.codecov.io/gh/scverse/spatialdata/pull/610?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) | Coverage Δ | |
|---|---|---|
| [src/spatialdata/\_core/spatialdata.py](https://app.codecov.io/gh/scverse/spatialdata/pull/610?src=pr&el=tree&filepath=src%2Fspatialdata%2F_core%2Fspatialdata.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NwYXRpYWxkYXRhL19jb3JlL3NwYXRpYWxkYXRhLnB5) | `91.01% <100.00%> (+0.11%)` | :arrow_up: |
| [src/spatialdata/models/\_utils.py](https://app.codecov.io/gh/scverse/spatialdata/pull/610?src=pr&el=tree&filepath=src%2Fspatialdata%2Fmodels%2F_utils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NwYXRpYWxkYXRhL21vZGVscy9fdXRpbHMucHk=) | `91.55% <100.00%> (+0.64%)` | :arrow_up: |
| [src/spatialdata/models/models.py](https://app.codecov.io/gh/scverse/spatialdata/pull/610?src=pr&el=tree&filepath=src%2Fspatialdata%2Fmodels%2Fmodels.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c3JjL3NwYXRpYWxkYXRhL21vZGVscy9tb2RlbHMucHk=) | `87.61% <100.00%> (-0.09%)` | :arrow_down: |

... and [4 files with indirect coverage changes](https://app.codecov.io/gh/scverse/spatialdata/pull/610/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse)

</details>

Patch:
diff --git a/src/spatialdata/_core/spatialdata.py b/src/spatialdata/_core/spatialdata.py
index 4be04f8..456bcba 100644
--- a/src/spatialdata/_core/spatialdata.py
+++ b/src/spatialdata/_core/spatialdata.py
@@ -37,7 +37,7 @@ from spatialdata.models import (
     get_model,
     get_table_keys,
 )
-from spatialdata.models._utils import SpatialElement, get_axes_names
+from spatialdata.models._utils import SpatialElement, convert_region_column_to_categorical, get_axes_names
 
 if TYPE_CHECKING:
     from spatialdata._core.query.spatial_query import BaseSpatialRequest
@@ -471,6 +471,7 @@ class SpatialData:
             self._set_table_annotation_target(table, region, region_key, instance_key)
         else:
             raise TypeError("No current annotation metadata found. Please specify both region_key and instance_key.")
+        convert_region_column_to_categorical(table)
 
     @property
     def query(self) -> QueryManager:
diff --git a/src/spatialdata/models/_utils.py b/src/spatialdata/models/_utils.py
index 55b4251..71cf8c9 100644
--- a/src/spatialdata/models/_utils.py
+++ b/src/spatialdata/models/_utils.py
@@ -1,11 +1,14 @@
 from __future__ import annotations
 
+import warnings
 from functools import singledispatch
 from typing import TYPE_CHECKING, Any, Union
 
 import dask.dataframe as dd
 import geopandas
 import numpy as np
+import pandas as pd
+from anndata import AnnData
 from dask.dataframe import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
@@ -359,3 +362,18 @@ def get_raster_model_from_data_dims(dims: tuple[str, ...]) -> type[RasterSchema]
     if C in dims:
         return Image3DModel if Z in dims else Image2DModel
     return Labels3DModel if Z in dims else Labels2DModel
+
+
+def convert_region_column_to_categorical(table: AnnData) -> AnnData:
+    from spatialdata.models.models import TableModel
+
+    if TableModel.ATTRS_KEY in table.uns:
+        region_key = table.uns[TableModel.ATTRS_KEY][TableModel.REGION_KEY_KEY]
+        if not isinstance(table.obs[region_key].dtype, pd.CategoricalDtype):
+            warnings.warn(
+                f"Converting `{TableModel.REGION_KEY_KEY}: {region_key}` to categorical dtype.",
+                UserWarning,
+                stacklevel=2,
+            )
+            table.obs[region_key] = pd.Categorical(table.obs[region_key])
+    return table
diff --git a/src/spatialdata/models/models.py b/src/spatialdata/models/models.py
index 943b669..bdf1c9e 100644
--- a/src/spatialdata/models/models.py
+++ b/src/spatialdata/models/models.py
@@ -43,6 +43,7 @@ from spatialdata.models._utils import (
     MappingToCoordinateSystem_t,
     SpatialElement,
     _validate_mapping_to_coordinate_system_type,
+    convert_region_column_to_categorical,
 )
 from spatialdata.transformations._utils import (
     _get_transformations,
@@ -939,11 +940,7 @@ class TableModel:
         region_: list[str] = region if isinstance(region, list) else [region]
         if not adata.obs[region_key].isin(region_).all():
             raise ValueError(f"`adata.obs[{region_key}]` values do not match with `{cls.REGION_KEY}` values.")
-        if not isinstance(adata.obs[region_key].dtype, CategoricalDtype):
-            warnings.warn(
-                f"Converting `{cls.REGION_KEY_KEY}: {region_key}` to categorical dtype.", UserWarning, stacklevel=2
-            )
-            adata.obs[region_key] = pd.Categorical(adata.obs[region_key])
+
         if instance_key is None:
             raise ValueError("`instance_key` must be provided.")
 
@@ -959,7 +956,7 @@ class TableModel:
         attr = {"region": region, "region_key": region_key, "instance_key": instance_key}
         adata.uns[cls.ATTRS_KEY] = attr
         cls().validate(adata)
-        return adata
+        return convert_region_column_to_categorical(adata)
 
 
 Schema_t = Union[


Test patch:
diff --git a/tests/io/test_multi_table.py b/tests/io/test_multi_table.py
index 8f6d818..4140633 100644
--- a/tests/io/test_multi_table.py
+++ b/tests/io/test_multi_table.py
@@ -95,7 +95,8 @@ class TestMultiTable:
         ):
             full_sdata.set_table_annotates_spatialelement("table", "non_existing")
 
-    def test_set_table_annotates_spatialelement(self, full_sdata):
+    def test_set_table_annotates_spatialelement(self, full_sdata, tmp_path):
+        tmpdir = Path(tmp_path) / "tmp.zarr"
         del full_sdata["table"].uns[TableModel.ATTRS_KEY]
         with pytest.raises(
             TypeError, match="No current annotation metadata found. " "Please specify both region_key and instance_key."
@@ -120,6 +121,12 @@ class TestMultiTable:
             "table", pd.Series(["circles", "poly"]), region_key="region", instance_key="instance_id"
         )
 
+        full_sdata["table"].obs["region"] = "circles"
+        full_sdata.set_table_annotates_spatialelement(
+            "table", "circles", region_key="region", instance_key="instance_id"
+        )
+        full_sdata.write(tmpdir)
+
     def test_old_accessor_deprecation(self, full_sdata, tmp_path):
         # To test self._backed
         tmpdir = Path(tmp_path) / "tmp.zarr"