Test new helper functions

Author: znicholls

src/pandas_openscm/accessors.py

@@ -30,6 +30,7 @@
 )
 from pandas_openscm.index_manipulation import (
     convert_index_to_category_index,
+    ensure_index_is_multiindex,
     set_index_levels_func,
     update_index_levels_from_other_func,
     update_index_levels_func,
@@ -74,6 +75,46 @@ def __init__(self, pandas_obj: pd.DataFrame):
         # However, it's probably better to do validation closer to the data use.
         self._df = pandas_obj
 
+    def ensure_index_is_multiindex(self, copy: bool = True) -> pd.DataFrame:
+        """
+        Ensure that the index is a [pd.MultiIndex][pandas.MultiIndex]
+
+        Parameters
+        ----------
+        copy
+            Whether to copy `df` before manipulating the index name
+
+        Returns
+        -------
+        :
+            `df` with 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 ensure_index_is_multiindex(self._df, copy=copy)
+
+    def eiim(self, copy: bool = True) -> pd.DataFrame:
+        """
+        Ensure that the index is a [pd.MultiIndex][pandas.MultiIndex]
+
+        Alias for [ensure_index_is_multiindex][(c).]
+
+        Parameters
+        ----------
+        copy
+            Whether to copy `df` before manipulating the index name
+
+        Returns
+        -------
+        :
+            `df` with 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_index_is_multiindex(copy=copy)
+
     def fix_index_name_after_groupby_quantile(
         self, new_name: str = "quantile", copy: bool = False
     ) -> pd.DataFrame:
@@ -496,6 +537,33 @@ def plot_plume_after_calculating_quantiles(  # noqa: PLR0913
             observed=observed,
         )
 
+    def set_index_levels(
+        self,
+        levels_to_set: dict[str, Any | Collection[Any]],
+        copy: bool = True,
+    ) -> pd.DataFrame:
+        """
+        Set the index levels
+
+        Parameters
+        ----------
+        levels_to_set
+            Mapping of level names to values to set
+
+        copy
+            Should the [pd.DataFrame][pandas.DataFrame] be copied before returning?
+
+        Returns
+        -------
+        :
+            [pd.DataFrame][pandas.DataFrame] with updates applied to its index
+        """
+        return set_index_levels_func(
+            self._df,
+            levels_to_set=levels_to_set,
+            copy=copy,
+        )
+
     def to_category_index(self) -> pd.DataFrame:
         """
         Convert the index's values to categories
@@ -717,33 +785,6 @@ def update_index_levels_from_other(
             remove_unused_levels=remove_unused_levels,
         )
 
-    def set_index_levels(
-        self,
-        levels_to_set: dict[str, Any | Collection[Any]],
-        copy: bool = True,
-    ) -> pd.DataFrame:
-        """
-        Set the index levels
-
-        Parameters
-        ----------
-        levels_to_set
-            Mapping of level names to values to set
-
-        copy
-            Should the [pd.DataFrame][pandas.DataFrame] be copied before returning?
-
-        Returns
-        -------
-        :
-            [pd.DataFrame][pandas.DataFrame] with updates applied to its index
-        """
-        return set_index_levels_func(
-            self._df,
-            levels_to_set=levels_to_set,
-            copy=copy,
-        )
-
 
 def register_pandas_accessor(namespace: str = "openscm") -> None:
     """

src/pandas_openscm/index_manipulation.py

@@ -90,14 +90,16 @@ def ensure_index_is_multiindex(pandas_obj: P, copy: bool = True) -> P:
     -------
     :
         `pandas_obj` with a [pd.MultiIndex][pandas.MultiIndex]
-    """
-    # TODOO: accessor and tests
-    if isinstance(pandas_obj.index, pd.MultiIndex):
-        return pandas_obj
 
+        If the index was already a [pd.MultiIndex][pandas.MultiIndex],
+        this is a no-op (although the value of copy is respected).
+    """
     if copy:
         pandas_obj = pandas_obj.copy()
 
+    if isinstance(pandas_obj.index, pd.MultiIndex):
+        return pandas_obj
+
     pandas_obj.index = ensure_is_multiindex(pandas_obj.index)
 
     return pandas_obj

tests/integration/index_manipulation/test_integration_index_manipulation_ensure_index_is_multiindex.py

@@ -5,40 +5,104 @@
 from __future__ import annotations
 
 import numpy as np
+import pandas as pd
+import pytest
 
-from pandas_openscm.index_manipulation import convert_index_to_category_index
+from pandas_openscm.index_manipulation import ensure_index_is_multiindex
 from pandas_openscm.testing import create_test_df
 
 
-def ensure_index_is_multiindex():
-    raise NotImplementedError
-    units = ["Mt", "kg", "W"]
+@pytest.mark.parametrize("copy, copy_exp", ((None, True), (True, True), (False, False)))
+def test_ensure_index_is_multiindex(copy, copy_exp):
+    start = pd.DataFrame(
+        [[1, 2], [3, 4]],
+        columns=[10, 20],
+        index=pd.Index(["a", "b"], name="variable"),
+    )
+
+    call_kwargs = {}
+    if copy is not None:
+        call_kwargs["copy"] = copy
+
+    res = ensure_index_is_multiindex(start, **call_kwargs)
+
+    assert isinstance(res.index, pd.MultiIndex)
+    assert res.index.names == ["variable"]
+
+    if copy_exp:
+        # New object returned
+        assert id(start) != id(res)
 
-    # Biggish DataFrame
+        assert isinstance(start.index, pd.Index)
+
+    else:
+        # Same object returned
+        assert id(start) == id(res)
+        # Therefore affects input object too
+        assert isinstance(start.index, pd.MultiIndex)
+
+
+@pytest.mark.parametrize("copy, copy_exp", ((None, True), (True, True), (False, False)))
+def test_ensure_index_is_multiindex_no_op(copy, copy_exp):
     start = create_test_df(
-        variables=[(f"variable_{i}", units[i % len(units)]) for i in range(25)],
-        n_scenarios=30,
-        n_runs=60,
+        variables=[("Temperature", "K")],
+        n_scenarios=2,
+        n_runs=2,
         timepoints=np.arange(1750.0, 2100.0 + 1.0),
     )
 
-    res = convert_index_to_category_index(start)
+    call_kwargs = {}
+    if copy is not None:
+        call_kwargs["copy"] = copy
 
-    run_checks(res, start)
+    res = ensure_index_is_multiindex(start, **call_kwargs)
 
+    # Already a MultiIndex, should be no change
+    pd.testing.assert_index_equal(res.index, start.index)
 
-def test_accessor(setup_pandas_accessor):
-    raise NotImplementedError
-    units = ["Mt", "kg", "W"]
+    # Behaviour of copy should be respected to avoid confusing behaviour
+    if copy_exp:
+        # New object returned
+        assert id(start) != id(res)
 
-    # Biggish DataFrame
-    start = create_test_df(
-        variables=[(f"variable_{i}", units[i % len(units)]) for i in range(25)],
-        n_scenarios=30,
-        n_runs=60,
-        timepoints=np.arange(1750.0, 2100.0 + 1.0),
+    else:
+        # Same object returned
+        assert id(start) == id(res)
+
+
+@pytest.mark.parametrize("copy, copy_exp", ((None, True), (True, True), (False, False)))
+def test_accessor(setup_pandas_accessor, copy, copy_exp):
+    start = pd.DataFrame(
+        [[1, 2], [3, 4]],
+        columns=[10, 20],
+        index=pd.Index(["a", "b"], name="variable"),
     )
 
-    res = start.openscm.to_category_index()
+    call_kwargs = {}
+    if copy is not None:
+        call_kwargs["copy"] = copy
+
+    res = start.openscm.ensure_index_is_multiindex(**call_kwargs)
+
+    assert isinstance(res.index, pd.MultiIndex)
+    assert res.index.names == ["variable"]
+    if copy_exp:
+        # New object returned
+        assert id(start) != id(res)
+
+    else:
+        # Same object returned
+        assert id(start) == id(res)
+
+    # Test alias too
+    res_short = start.openscm.eiim(**call_kwargs)
+
+    assert isinstance(res_short.index, pd.MultiIndex)
+    assert res_short.index.names == ["variable"]
+    if copy_exp:
+        # New object returned
+        assert id(start) != id(res_short)
 
-    run_checks(res, start)
+    else:
+        # Same object returned
+        assert id(start) == id(res_short)