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..26e88b8 --- /dev/null +++ b/xarray_ms/accessors.py @@ -0,0 +1,21 @@ +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] diff --git a/xarray_ms/backend/msv2/entrypoint.py b/xarray_ms/backend/msv2/entrypoint.py index 35f8ee9..f172545 100644 --- a/xarray_ms/backend/msv2/entrypoint.py +++ b/xarray_ms/backend/msv2/entrypoint.py @@ -375,6 +375,7 @@ def open_datatree( antenna_factory = AntennaDatasetFactory(structure_factory) key = ",".join(f"{k}={v}" for k, v in sorted(partition_key)) + ds.attrs["antenna_xds_link"] = f"{key}/ANTENNA" datasets[key] = ds datasets[f"{key}/ANTENNA"] = antenna_factory.get_dataset()