Add accessor for set levels

Author: znicholls

src/pandas_openscm/accessors/series.py

@@ -4,11 +4,14 @@
 
 from __future__ import annotations
 
-from collections.abc import Mapping
+from collections.abc import Collection, Mapping
 from typing import TYPE_CHECKING, Any, Generic, TypeVar
 
 import pandas as pd
 
+from pandas_openscm.index_manipulation import (
+    set_index_levels_func,
+)
 from pandas_openscm.unit_conversion import convert_unit, convert_unit_like
 
 if TYPE_CHECKING:
@@ -305,33 +308,33 @@ def convert_unit_like(
     #     """
     #     return mi_loc(self._df, locator)
 
-    # 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 set_index_levels(
+        self,
+        levels_to_set: dict[str, Any | Collection[Any]],
+        copy: bool = True,
+    ) -> S:
+        """
+        Set the index levels
+
+        Parameters
+        ----------
+        levels_to_set
+            Mapping of level names to values to set
+
+        copy
+            Should the [pd.Series][pandas.Series] be copied before returning?
+
+        Returns
+        -------
+        :
+            [pd.Series][pandas.Series] with updates applied to its index
+        """
+        return set_index_levels_func(
+            self._series,
+            levels_to_set=levels_to_set,
+            copy=copy,
+        )
+
     # def to_category_index(self) -> pd.DataFrame:
     #     """
     #     Convert the index's values to categories

src/pandas_openscm/index_manipulation.py

@@ -910,7 +910,6 @@ def set_levels(
 
 
 def set_index_levels_func(
-    # TODO: check support for series and add accessors
     pobj: P,
     levels_to_set: dict[str, Any | Collection[Any]],
     copy: bool = True,

src/pandas_openscm/testing.py

@@ -10,7 +10,7 @@
 
 import itertools
 from collections.abc import Collection
-from typing import TYPE_CHECKING, Any, cast
+from typing import TYPE_CHECKING, Any, Literal, TypeVar, cast, overload
 
 import numpy as np
 import pandas as pd
@@ -22,6 +22,8 @@
 if TYPE_CHECKING:
     import pytest
 
+    P = TypeVar("P", pd.DataFrame | pd.Series[Any])
+
 
 def get_db_data_backends() -> tuple[type[object], ...]:
     return tuple(v[1] for v in DATA_BACKENDS.options)
@@ -97,6 +99,76 @@ def assert_frame_alike(
     )
 
 
+@overload
+def convert_to_desired_type(
+    pobj: pd.DataFrame, pobj_type: Literal["DataFrame"]
+) -> pd.DataFrame: ...
+
+
+@overload
+def convert_to_desired_type(
+    pobj: pd.DataFrame, pobj_type: Literal["Series"]
+) -> pd.Series[Any]: ...
+
+
+def convert_to_desired_type(
+    df: pd.DataFrame, pobj_type: Literal["DataFrame", "Series"]
+) -> pd.DataFrame | pd.Series[Any]:
+    """
+    Convert a `df` to the desired type for testing
+
+    Parameters
+    ----------
+    df
+        [pd.DataFrame][pandas.DataFrame] to convert
+
+    pobj_type
+        Type to convert to
+
+        If "DataFrame", then `df` is simply returned.
+        If "Series", then the first column of `df` is returned.
+
+    Returns
+    -------
+    :
+        `df` converted to the desired type
+    """
+    if pobj_type == "DataFrame":
+        return df
+
+    if pobj_type == "Series":
+        res = df[df.columns[0]]
+        return res
+
+    raise NotImplementedError(pobj_type)
+
+
+def check_result(res: P, exp: P) -> None:
+    """
+    Check result in the case where it could be multiple types
+
+    Specifically, [pd.DataFrame][pandas.DataFrame]
+    or [pd.Series][pandas.Series].
+
+    This is a thin wrapper, if you want specific functionality,
+    use the underlying function.
+
+    Parameters
+    ----------
+    res
+        Result
+
+    exp
+        Expected
+    """
+    if isinstance(res, pd.DataFrame):
+        assert_frame_alike(res, exp)
+    elif isinstance(res, pd.Series):
+        pd.testing.assert_series_equal(res, exp)
+    else:
+        raise NotImplementedError(type(res))
+
+
 def create_test_df(
     *,
     variables: Collection[tuple[str, str]],

tests/integration/index_manipulation/test_integration_index_manipulation_set_levels.py

@@ -9,6 +9,15 @@
 import pytest
 
 from pandas_openscm.index_manipulation import set_index_levels_func, set_levels
+from pandas_openscm.testing import convert_to_desired_type
+
+pobj_type = pytest.mark.parametrize(
+    "pobj_type",
+    ("DataFrame", "Series"),
+)
+"""
+Parameterisation to use to check handling of both DataFrame and Series
+"""
 
 
 @pytest.mark.parametrize(
@@ -160,7 +169,8 @@ def test_set_levels(start, levels_to_set, exp):
     pd.testing.assert_index_equal(res, exp)
 
 
-def test_set_levels_with_a_dataframe():
+@pobj_type
+def test_set_levels_with_a_dataframe(pobj_type):
     start = pd.MultiIndex.from_tuples(
         [
             ("sa", "va", "kg", 0),
@@ -170,11 +180,14 @@ def test_set_levels_with_a_dataframe():
         ],
         names=["scenario", "variable", "unit", "run_id"],
     )
-    start_df = pd.DataFrame(
-        np.zeros((start.shape[0], 3)), columns=[2010, 2020, 2030], index=start
+    start_pobj = convert_to_desired_type(
+        pd.DataFrame(
+            np.zeros((start.shape[0], 3)), columns=[2010, 2020, 2030], index=start
+        ),
+        pobj_type,
     )
 
-    res = set_index_levels_func(start_df, levels_to_set={"new_variable": "test"})
+    res = set_index_levels_func(start_pobj, levels_to_set={"new_variable": "test"})
 
     exp = pd.MultiIndex.from_tuples(
         [
@@ -189,11 +202,13 @@ def test_set_levels_with_a_dataframe():
     pd.testing.assert_index_equal(res.index, exp)
 
 
-def test_set_levels_raises_type_error():
+@pobj_type
+def test_set_levels_raises_type_error(pobj_type):
     start = pd.DataFrame(
         np.arange(2 * 4).reshape((4, 2)),
         columns=[2010, 2020],
     )
+    start = convert_to_desired_type(start, pobj_type)
 
     levels_to_set = {"new_variable": "test"}
 
@@ -221,7 +236,7 @@ def test_set_levels_raises_value_error():
         set_levels(start, levels_to_set=levels_to_set)
 
 
-def test_accessor(setup_pandas_accessors):
+def test_accessor_df(setup_pandas_accessors):
     start = pd.DataFrame(
         np.arange(2 * 4).reshape((4, 2)),
         columns=[2010, 2020],
@@ -262,3 +277,44 @@ def test_accessor(setup_pandas_accessors):
     # Test function too
     res = set_index_levels_func(start, levels_to_set=levels_to_set)
     pd.testing.assert_frame_equal(res, exp)
+
+
+def test_accessor_series(setup_pandas_accessors):
+    start = pd.Series(
+        np.arange(4),
+        index=pd.MultiIndex.from_tuples(
+            [
+                ("sa", "va", "kg", 0),
+                ("sb", "vb", "m", -1),
+                ("sa", "va", "kg", -2),
+                ("sa", "vb", "kg", 2),
+            ],
+            names=["scenario", "variable", "unit", "run_id"],
+        ),
+    )
+
+    levels_to_set = {
+        "model_id": "674",
+        "unit": ["t", "km", "g", "kg"],
+        "scenario": 1,
+    }
+
+    exp = pd.Series(
+        start.values,
+        index=pd.MultiIndex.from_tuples(
+            [
+                (1, "va", "t", 0, "674"),
+                (1, "vb", "km", -1, "674"),
+                (1, "va", "g", -2, "674"),
+                (1, "vb", "kg", 2, "674"),
+            ],
+            names=["scenario", "variable", "unit", "run_id", "model_id"],
+        ),
+    )
+
+    res = start.openscm.set_index_levels(levels_to_set=levels_to_set)
+    pd.testing.assert_series_equal(res, exp)
+
+    # Test function too
+    res = set_index_levels_func(start, levels_to_set=levels_to_set)
+    pd.testing.assert_series_equal(res, exp)

tests/integration/test_unit_conversion.py

@@ -6,7 +6,7 @@
 
 import re
 import sys
-from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload
+from typing import TYPE_CHECKING, Any, TypeVar
 from unittest.mock import patch
 
 import numpy as np
@@ -17,7 +17,12 @@
 from pandas_openscm.index_manipulation import (
     set_index_levels_func,
 )
-from pandas_openscm.testing import assert_frame_alike, create_test_df
+from pandas_openscm.testing import (
+    assert_frame_alike,
+    check_result,
+    convert_to_desired_type,
+    create_test_df,
+)
 from pandas_openscm.unit_conversion import (
     AmbiguousTargetUnitError,
     MissingDesiredUnitError,
@@ -29,6 +34,14 @@
 if TYPE_CHECKING:
     P = TypeVar("P", pd.DataFrame | pd.Series[Any])
 
+pobj_type = pytest.mark.parametrize(
+    "pobj_type",
+    ("DataFrame", "Series"),
+)
+"""
+Parameterisation to use to check handling of both DataFrame and Series
+"""
+
 check_auto_index_casting_pobj = pytest.mark.parametrize(
     "only_two_index_levels_pobj",
     (
@@ -43,48 +56,6 @@
 This parameterisation ensures that we check this edge case.
 """
 
-pobj_type = pytest.mark.parametrize(
-    "pobj_type",
-    ("DataFrame", "Series"),
-)
-"""
-Parameterisation to use to check handling of both DataFrame and Series
-"""
-
-
-@overload
-def convert_to_desired_type(
-    pobj: pd.DataFrame, pobj_type: Literal["DataFrame"]
-) -> pd.DataFrame: ...
-
-
-@overload
-def convert_to_desired_type(
-    pobj: pd.DataFrame, pobj_type: Literal["Series"]
-) -> pd.Series[Any]: ...
-
-
-def convert_to_desired_type(
-    df: pd.DataFrame, pobj_type: Literal["DataFrame", "Series"]
-) -> pd.DataFrame | pd.Series[Any]:
-    if pobj_type == "DataFrame":
-        return df
-
-    if pobj_type == "Series":
-        res = df[df.columns[0]]
-        return res
-
-    raise NotImplementedError(pobj_type)
-
-
-def check_result(res: P, exp: P) -> None:
-    if isinstance(res, pd.DataFrame):
-        assert_frame_alike(res, exp)
-    elif isinstance(res, pd.Series):
-        pd.testing.assert_series_equal(res, exp)
-    else:
-        raise NotImplementedError(type(res))
-
 
 @pobj_type
 @check_auto_index_casting_pobj