diff --git a/xarray_ms/__init__.py b/xarray_ms/__init__.py index e69de29..d4bd2a8 100644 --- a/xarray_ms/__init__.py +++ b/xarray_ms/__init__.py @@ -0,0 +1 @@ +import xarray_ms.accessors # noqa: F401 diff --git a/xarray_ms/accessors.py b/xarray_ms/accessors.py new file mode 100644 index 0000000..90968f4 --- /dev/null +++ b/xarray_ms/accessors.py @@ -0,0 +1,31 @@ +from typing import Any + +from xarray.core.datatree import DataTree +from xarray.core.extensions import ( + register_datatree_accessor, +) + + +@register_datatree_accessor("subtable") +class SubTableAccessor: + def __init__(self, node: DataTree): + self.node = node + + @property + def antenna(self) -> DataTree: + """Returns the antenna dataset""" + + try: + link = self.node.attrs["antenna_xds_link"] + except KeyError: + raise ValueError("antenna_xds_link not found") + else: + return self.node.root[link] + + def __getitem__(self, key: Any) -> Any: + try: + link = self.node.attrs[key] + except KeyError: + raise ValueError(f"{key} link attribute not found") + else: + return self.node.root[link] diff --git a/xarray_ms/backend/msv2/entrypoint.py b/xarray_ms/backend/msv2/entrypoint.py index 8b0d037..3be2878 100644 --- a/xarray_ms/backend/msv2/entrypoint.py +++ b/xarray_ms/backend/msv2/entrypoint.py @@ -439,6 +439,7 @@ def open_groups_as_dict( antenna_factory = AntennaDatasetFactory(structure_factory) path = f"{ms_name}/partition_{p:03}" + ds.attrs["antenna_xds_link"] = f"{path}/ANTENNA" datasets[path] = ds datasets[f"{path}/ANTENNA"] = antenna_factory.get_dataset()