general fixes to adapt table model

src/spatialdata/_core/query/spatial_query.py

@@ -8,6 +8,7 @@
 import dask.array as da
 import numpy as np
 import pandas as pd
+from anndata import AnnData
 from dask.dataframe.core import DataFrame as DaskDataFrame
 from datatree import DataTree
 from geopandas import GeoDataFrame
@@ -17,6 +18,7 @@
 from tqdm import tqdm
 from xarray import DataArray
 
+from spatialdata._core._elements import Tables
 from spatialdata._core.query._utils import get_bounding_box_corners
 from spatialdata._core.spatialdata import SpatialData
 from spatialdata._logging import logger
@@ -266,7 +268,9 @@ def _(
         new_elements[element_type] = queried_elements
 
     if filter_table:
-        tables = {name: _filter_table_by_elements(table, new_elements) for name, table in sdata.tables.items()}
+        tables: dict[str, AnnData] | Tables = {
+            name: _filter_table_by_elements(table, new_elements) for name, table in sdata.tables.items()
+        }
     else:
         tables = sdata.tables
 
@@ -645,7 +649,9 @@ def _polygon_query(
 
     if filter_table:
         elements = {"shapes": new_shapes, "points": new_points}
-        tables = {name: _filter_table_by_elements(table, elements) for name, table in sdata.tables.items()}
+        tables: dict[str, AnnData] | Tables = {
+            name: _filter_table_by_elements(table, elements) for name, table in sdata.tables.items()
+        }
     else:
         tables = sdata.tables
     return SpatialData(shapes=new_shapes, points=new_points, images=new_images, tables=tables)
@@ -754,7 +760,9 @@ def polygon_query(
         geodataframes[k] = vv
     if filter_table:
         elements = {"shapes": geodataframes}
-        tables = {name: _filter_table_by_elements(table, elements) for name, table in sdata.tables.items()}
+        tables: dict[str, AnnData] | Tables = {
+            name: _filter_table_by_elements(table, elements) for name, table in sdata.tables.items()
+        }
     else:
         tables = sdata.tables
 

src/spatialdata/_core/spatialdata.py

@@ -167,9 +167,8 @@ def __init__(
                 self.points[k] = v
 
         if tables is not None:
-            self._tables: dict[str, AnnData] = {}
-
-            self._add_tables(table_mapping=tables)
+            for k, v in tables.items():
+                self.tables[k] = v
 
         self._query = QueryManager(self)
 
@@ -672,7 +671,7 @@ def filter_by_coordinate_system(
         # TODO: check whether full table dict should be returned or only those which annotate elements. Also check
         # filtering with tables having potentially different keys.
         if filter_tables:
-            tables = {}
+            tables: Tables = Tables(shared_keys=self._shared_keys)
             for table_name, table in self._tables.items():
                 tables[table_name] = _filter_table_by_coordinate_system(table, element_paths_in_coordinate_system)
         else:
@@ -975,7 +974,7 @@ def write(
         assert isinstance(self.path, Path)
 
     @property
-    def tables(self) -> dict[str, AnnData]:
+    def tables(self) -> Tables:
         """
         Return tables dictionary.
 
@@ -987,6 +986,14 @@ def tables(self) -> dict[str, AnnData]:
         """
         return self._tables
 
+    @tables.setter
+    def tables(self, shapes: dict[str, GeoDataFrame]) -> None:
+        """Set shapes."""
+        self._shared_keys = self._shared_keys - set(self._tables.keys())
+        self._tables = Tables(shared_keys=self._shared_keys)
+        for k, v in shapes.items():
+            self._tables[k] = v
+
     @property
     def table(self) -> None | AnnData:
         """
@@ -1003,8 +1010,8 @@ def table(self) -> None | AnnData:
             stacklevel=2,
         )
         # Isinstance will still return table if anndata has 0 rows.
-        if isinstance(self._tables.get("table"), AnnData):
-            return self._tables["table"]
+        if isinstance(self.tables.get("table"), AnnData):
+            return self.tables["table"]
         return None
 
     @table.setter
@@ -1015,10 +1022,9 @@ def table(self, table: AnnData) -> None:
             stacklevel=2,
         )
         TableModel().validate(table)
-        if self._tables.get("table") is not None:
+        if self.tables.get("table") is not None:
             raise ValueError("The table already exists. Use del sdata.tables['table'] to remove it first.")
-        self._tables["table"] = table
-        self._store_tables(table_name="table")
+        self.tables["table"] = table
 
     @table.deleter
     def table(self) -> None:
@@ -1038,106 +1044,6 @@ def table(self) -> None:
             # More informative than the error in the zarr library.
             raise KeyError("table with name 'table' not present in the SpatialData object.")
 
-    def add_tables(
-        self,
-        table: None | AnnData = None,
-        table_name: None | str = None,
-        table_mapping: None | dict[str, AnnData] = None,
-    ) -> None:
-        """
-        Add AnnData tables to the SpatialData object.
-
-        Parameters
-        ----------
-        table : None or AnnData, optional
-            The table to be added. If provided, it is added with the specified table_name.
-            If not provided, the table_mapping must be provided.
-
-        table_name : None or str, optional
-            The name of the table to be added. This parameter is required if table is provided.
-
-        table_mapping : None or dict[str, AnnData], optional
-            A dictionary mapping table names to tables. If provided, the tables are added
-            according to the mapping. This parameter is required if table is not provided.
-        """
-        self._add_tables(table=table, table_name=table_name, table_mapping=table_mapping)
-
-    def _store_tables(self, table_name: None | str = None, table_mapping: None | dict[str, AnnData] = None) -> None:
-        """
-        Write tables back to the SpatialData Zarr store.
-
-        This method is used to store tables in a backed `AnnData` object. If the `AnnData` object is backed, the tables
-        are stored in a Zarr store, under the "tables" group.
-        attribute.
-
-        Parameters
-        ----------
-        table_name : None or str, optional
-            The name of the table to store. If not provided, all tables in `table_mapping` will be stored.
-        table_mapping : None or dict[str, AnnData], optional
-            A dictionary mapping table names to `AnnData` tables. If provided, all tables in `table_mapping` will be
-            stored.
-
-        Raises
-        ------
-        TypeError
-            If both `table_name` and `table_mapping` are None.
-        """
-        if self.is_backed():
-            from spatialdata._io.io_table import write_table
-
-            store = parse_url(self.path, mode="r+").store
-            root = zarr.group(store=store)
-            elem_group = root.require_group(name="tables")
-            if table_name:
-                write_table(table=self._tables[table_name], group=elem_group, name=table_name)
-            elif table_mapping:
-                for table_name in table_mapping:
-                    write_table(table=self._tables[table_name], group=elem_group, name=table_name)
-            else:
-                raise TypeError("Missing arguments, either table_name or table_mapping should be provided.")
-
-    def _add_tables(
-        self,
-        table: None | AnnData = None,
-        table_name: None | str = None,
-        table_mapping: None | dict[str, AnnData] = None,
-    ) -> None:
-        """
-        Validate and add AnnData table(s) to the SpatialData object and save to the SpatialData Zarr store if backed.
-
-        Parameters
-        ----------
-        table : None or AnnData, optional
-            An optional AnnData table to be added.
-        table_name : None or str, optional
-            An optional string representing the name of the table to be added.
-        table_mapping : None or dict[str, AnnData], optional
-            An optional dictionary mapping table names to AnnData tables. If specified, multiple tables are added at
-            once
-
-        Raises
-        ------
-        ValueError
-            If the specified table already exists in the SpatialData.tables.
-        ValueError
-            If table is given but table_name is not specified.
-        """
-        if table:
-            if table_name:
-                TableModel().validate(table)
-                if self._tables.get(table_name) is not None:
-                    raise ValueError("The table already exists. Use del sdata.tables[<table_name>] to remove it first.")
-                self._tables[table_name] = table
-                self._store_tables(table_name=table_name)
-            else:
-                raise ValueError("Please provide a string value for the parameter table_name.")
-        elif table_mapping:
-            for table in table_mapping.values():
-                self.validate_table_in_spatialdata(table)
-            self._tables.update(table_mapping)
-            self._store_tables(table_mapping=table_mapping)
-
     @staticmethod
     def read(file_path: Path | str, selection: tuple[str] | None = None) -> SpatialData:
         """
@@ -1497,7 +1403,7 @@ def __setitem__(self, key: str, value: SpatialElement | AnnData) -> None:
         elif schema == ShapesModel:
             self.shapes[key] = value
         elif schema == TableModel:
-            self.add_tables(table_name=key, table=value)
+            self.tables[key] = value
         else:
             raise TypeError(f"Unknown element type with schema: {schema!r}.")