From 25f132b9063f6f98fe4570ff15958523294920b8 Mon Sep 17 00:00:00 2001
From: Simon Perkins <simon.perkins@gmail.com>
Date: Tue, 15 Oct 2024 09:34:42 +0200
Subject: [PATCH 1/2] Experiment with accessors and links in attributes

---
 xarray_ms/__init__.py                |  1 +
 xarray_ms/accessors.py               | 21 +++++++++++++++++++++
 xarray_ms/backend/msv2/entrypoint.py |  1 +
 3 files changed, 23 insertions(+)
 create mode 100644 xarray_ms/accessors.py

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()
 

From 57e4816ab6beca9070e3ec7424291346b4f2756a Mon Sep 17 00:00:00 2001
From: Simon Perkins <simon.perkins@gmail.com>
Date: Fri, 17 Jan 2025 08:37:45 +0200
Subject: [PATCH 2/2] Test __getitem__ idea

---
 xarray_ms/accessors.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/xarray_ms/accessors.py b/xarray_ms/accessors.py
index 26e88b8..90968f4 100644
--- a/xarray_ms/accessors.py
+++ b/xarray_ms/accessors.py
@@ -1,3 +1,5 @@
+from typing import Any
+
 from xarray.core.datatree import DataTree
 from xarray.core.extensions import (
   register_datatree_accessor,
@@ -19,3 +21,11 @@ def antenna(self) -> DataTree:
       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]