clean up

Author: crdanielbusch

src/pandas_openscm/index_manipulation.py

@@ -728,7 +728,9 @@ def set_levels(
         Input MultiIndex
 
     levels_to_set
-        Mapping of level names to values to set
+        Mapping of level names to values to set. If values is of type `Collection`,
+        it must be of the same length as the MultiIndex. If it is not a `Collection`,
+        it will be set to the same value for all levels.
 
     Returns
     -------
@@ -742,13 +744,15 @@ def set_levels(
         ValueError
             If the length of the values is not equal to the length of the index
     """
+    # TODO mypy says this is unreachable
     if not isinstance(ini, pd.MultiIndex):
         msg = f"Expected MultiIndex, got {type(ini)}"
         raise TypeError(msg)
 
     df = ini.to_frame(index=False)
 
     for level, value in levels_to_set.items():
+        # TODO do we need the isinstance check for strings here?
         if isinstance(value, Collection) and not isinstance(value, str):
             if len(value) != len(ini):
                 msg = (

tests/integration/index_manipulation/test_integration_index_manipulation_set_levels.py

@@ -159,9 +159,6 @@ def test_set_levels_with_a_dataframe():
 
 
 def test_set_levels_raises_type_error():
-    """
-    Test that set_levels raises an error when the levels to set are not in the index.
-    """
     start = pd.Index(range(10), name="index")
 
     levels_to_set = {"new_variable": "test"}
@@ -171,10 +168,6 @@ def test_set_levels_raises_type_error():
 
 
 def test_set_levels_raises_value_error():
-    """
-    Test that set_levels raises an error when the levels to set are same length
-    as index.
-    """
     start = pd.MultiIndex.from_tuples(
         [
             ("sa", "va", "kg", 0),