diff --git a/src/spatialdata/_core/spatialdata.py b/src/spatialdata/_core/spatialdata.py index db9b91ab..d05dac85 100644 --- a/src/spatialdata/_core/spatialdata.py +++ b/src/spatialdata/_core/spatialdata.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import json import os import warnings from collections.abc import Generator, Mapping @@ -1642,6 +1643,78 @@ def write_metadata( if consolidate_metadata: self.write_consolidated_metadata() + def get_attrs( + self, + key: str, + return_as: Literal["dict", "json", "df"] | None = None, + sep_for_nested_keys: str = "_", + flatten: bool = True, + ) -> dict[str, Any] | str | pd.DataFrame: + """ + Return the keys from sdata.attrs in a specific format. + + Parameters + ---------- + key + The key to retrieve from the attrs. + return_as + The format to return the data. Options are 'dict', 'json', 'df'. + If None, the function infers the format based on the data type. + sep_for_nested_keys : str, optional + Separator for nested keys in flattened data, by default "_". + flatten + If True, flatten the data if it is a mapping. + + """ + + def _flatten_mapping(m: Mapping[str, Any], parent_key: str = "", sep: str = "_") -> dict[str, Any]: + + items: list[tuple[str, Any]] = [] + for k, v in m.items(): + new_key = f"{parent_key}{sep}{k}" if parent_key else k + if isinstance(v, Mapping): + items.extend(_flatten_mapping(v, new_key, sep=sep).items()) + else: + items.append((new_key, v)) + return dict(items) + + if not isinstance(key, str): + raise TypeError("The key must be a string.") + + if not isinstance(sep_for_nested_keys, str): + raise TypeError("Parameter 'sep_for_nested_keys' must be a string.") + + if key not in self.attrs: + raise KeyError(f"The key '{key}' was not found in sdata.attrs.") + + data = self.attrs[key] + + # If the data is a mapping, flatten it + if flatten and isinstance(data, Mapping): + data = _flatten_mapping(data, sep=sep_for_nested_keys) + + if return_as is None: + return data + + if return_as == "dict": + if not isinstance(data, dict): + raise TypeError("Cannot convert non-dictionary data to a dictionary.") + return data + + if return_as == "json": + try: + return json.dumps(data) + except (TypeError, ValueError) as e: + raise ValueError(f"Failed to convert data to JSON: {e}") from e + + if return_as == "df": + try: + return pd.DataFrame([data]) + except Exception as e: + raise ValueError(f"Failed to convert data to DataFrame: {e}") from e + + raise ValueError(f"Invalid 'return_as' value: {return_as}. Expected 'dict', 'json', 'df', or None.") + @property def tables(self) -> Tables: """