Add index accessor

Author: znicholls

src/pandas_openscm/accessors/__init__.py

@@ -38,6 +38,7 @@
 import pandas as pd
 
 from pandas_openscm.accessors.dataframe import PandasDataFrameOpenSCMAccessor
+from pandas_openscm.accessors.index import PandasIndexOpenSCMAccessor
 from pandas_openscm.accessors.series import PandasSeriesOpenSCMAccessor
 
 
@@ -73,4 +74,4 @@ def register_pandas_accessors(namespace: str = "openscm") -> None:
         PandasDataFrameOpenSCMAccessor
     )
     pd.api.extensions.register_series_accessor(namespace)(PandasSeriesOpenSCMAccessor)
-    # pd.api.extensions.register_index_accessor(namespace)(PandasIndexOpenSCMAccessor)
+    pd.api.extensions.register_index_accessor(namespace)(PandasIndexOpenSCMAccessor)

src/pandas_openscm/accessors/index.py

@@ -0,0 +1,74 @@
+"""
+Accessor for [pd.Index][pandas.Index] (and sub-classes)
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Generic, TypeVar
+
+import pandas as pd
+
+from pandas_openscm.index_manipulation import ensure_is_multiindex
+
+if TYPE_CHECKING:
+    # Hmm this is somehow not correct.
+    # Figuring it out is a job for another day
+    Idx = TypeVar("Idx", bound=pd.Index[Any])
+
+
+else:
+    Idx = TypeVar("Idx")
+
+
+class PandasIndexOpenSCMAccessor(Generic[Idx]):
+    """
+    [pd.Index][pandas.Index] accessor
+
+    For details, see
+    [pandas' docs](https://pandas.pydata.org/docs/development/extending.html#registering-custom-accessors).
+    """
+
+    def __init__(self, index: Idx):
+        """
+        Initialise
+
+        Parameters
+        ----------
+        index
+            [pd.Index][pandas.Index] to use via the accessor
+        """
+        # It is possible to validate here.
+        # However, it's probably better to do validation closer to the data use.
+        self._index = index
+
+    def ensure_is_multiindex(self) -> pd.MultiIndex:
+        """
+        Ensure that the index is a [pd.MultiIndex][pandas.MultiIndex]
+
+        Returns
+        -------
+        :
+            `index` as a [pd.MultiIndex][pandas.MultiIndex]
+
+            If the index was already a [pd.MultiIndex][pandas.MultiIndex],
+            this is a no-op.
+        """
+        res = ensure_is_multiindex(self._index)
+
+        return res
+
+    def eim(self) -> pd.MultiIndex:
+        """
+        Ensure that the index is a [pd.MultiIndex][pandas.MultiIndex]
+
+        Alias for [ensure_is_multiindex][pandas_openscm.index_manipulation.]
+
+        Returns
+        -------
+        :
+            `index` as a [pd.MultiIndex][pandas.MultiIndex]
+
+            If the index was already a [pd.MultiIndex][pandas.MultiIndex],
+            this is a no-op (although the value of copy is respected).
+        """
+        return self.ensure_is_multiindex()

tests/conftest.py

@@ -45,3 +45,7 @@ def setup_pandas_accessors() -> None:
     pd.Series._accessors.discard("openscm")
     if hasattr(pd.Series, "openscm"):
         del pd.Series.openscm
+
+    pd.Index._accessors.discard("openscm")
+    if hasattr(pd.Index, "openscm"):
+        del pd.Index.openscm

tests/integration/index_manipulation/test_integration_index_manipulation_ensure_index_is_multiindex.py

@@ -178,6 +178,8 @@ def test_ensure_is_multiindex_accessor_index(setup_pandas_accessors):
     start = pd.Index([1, 2, 3], name="id")
 
     res = start.openscm.ensure_is_multiindex()
+    res_short = start.openscm.eim()
+    pd.testing.assert_index_equal(res, res_short)
 
     assert isinstance(res, pd.MultiIndex)
 
@@ -203,6 +205,8 @@ def test_ensure_is_multiindex_accessor_multiindex(setup_pandas_accessors):
     )
 
     res = start.openscm.ensure_is_multiindex()
+    res_short = start.openscm.eim()
+    pd.testing.assert_index_equal(res, res_short)
 
     # Same object returned
     assert id(start) == id(res)