Repository: pyjanitor-devs/pyjanitor (python)

Issue:
dtype conversion on index
# Brief Description

<!-- Please provide a brief description of what you'd like to propose. -->

I would like to propose a chainable method for index dtype conversion.

# Example API

```python
df.axis_astype(str, axis='index')
```


Pull request description:
[ENH] change dtype of index
# PR Description

Please describe the changes proposed in the pull request:

- change dtype of index/multiindex
- reuses the existing codes(for MultiIndex) for perf

**This PR resolves #1314 .**

<!-- As you go down the PR template, please feel free to delete sections that are irrelevant. -->

# PR Checklist

<!-- This checklist exists for newcomers who are not yet familiar with our requirements. If you are experienced with
the project, please feel free to delete this section. -->

Please ensure that you have done the following:

1. [x] PR in from a fork off your branch. Do not PR from `<your_username>`:`dev`, but rather from `<your_username>`:`<feature-branch_name>`.
<!-- Doing this helps us keep the commit history much cleaner than it would otherwise be. -->
2. [x] If you're not on the contributors list, add yourself to `AUTHORS.md`.
<!-- We'd like to acknowledge your contributions! -->
3. [x] Add a line to `CHANGELOG.md` under the latest version header (i.e. the one that is "on deck") describing the contribution.
    - Do use some discretion here; if there are multiple PRs that are related, keep them in a single line.

# Automatic checks

There will be automatic checks run on the PR. These include:

- Building a preview of the docs on Netlify
- Automatically linting the code
- Making sure the code is documented
- Making sure that all tests are passed
- Making sure that code coverage doesn't go down.

# Relevant Reviewers

<!-- Finally, please tag relevant maintainers to review. -->

Please tag maintainers to review.

- @ericmjl


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c6b874a..f03d56c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,8 @@
 -   [ENH] `conditional_join` now supports timedelta dtype. - PR #1297 @samukweku
 -   [ENH]  `get_join_indices` function added - returns only join indices between two dataframes. Issue #1310 @samukweku
 -   [ENH] `explode_index` function added. - Issue #1283
+-   [ENH] `conditional_join` now supports timedelta dtype. - PR #1297
+-   [ENH] `change_index_dtype` added. - @samukweku Issue #1314
 
 ## [v0.26.0] - 2023-09-18
 
diff --git a/janitor/functions/__init__.py b/janitor/functions/__init__.py
index 7897e50..9ee6bdb 100644
--- a/janitor/functions/__init__.py
+++ b/janitor/functions/__init__.py
@@ -20,6 +20,7 @@ from .add_columns import add_columns
 from .also import also
 from .bin_numeric import bin_numeric
 from .case_when import case_when
+from .change_index_dtype import change_index_dtype
 from .change_type import change_type
 from .clean_names import clean_names
 from .coalesce import coalesce
@@ -91,6 +92,7 @@ __all__ = [
     "bin_numeric",
     "case_when",
     "change_type",
+    "change_index_dtype",
     "clean_names",
     "coalesce",
     "collapse_levels",
diff --git a/janitor/functions/change_index_dtype.py b/janitor/functions/change_index_dtype.py
new file mode 100644
index 0000000..d7f6e61
--- /dev/null
+++ b/janitor/functions/change_index_dtype.py
@@ -0,0 +1,159 @@
+"""Implementation of the `change_index_dtype` function."""
+from __future__ import annotations
+
+from typing import Union
+
+import pandas as pd
+import pandas_flavor as pf
+
+from janitor.utils import check
+
+
+@pf.register_dataframe_method
+def change_index_dtype(
+    df: pd.DataFrame, dtype: Union[str, dict], axis: str = "index"
+) -> pd.DataFrame:
+    """Cast an index to a specified dtype ``dtype``.
+
+    This method does not mutate the original DataFrame.
+
+    Examples:
+        >>> import pandas as pd
+        >>> import numpy as np
+        >>> import janitor
+        >>> rng = np.random.default_rng(seed=0)
+        >>> np.random.seed(0)
+        >>> tuples = list(zip(*[['bar', 'bar', 'baz', 'baz',
+        ...             'foo', 'foo', 'qux', 'qux'],
+        ...              [1.0, 2.0, 1.0, 2.0,
+        ...               1.0, 2.0, 1.0, 2.0]]))
+        >>> idx = pd.MultiIndex.from_tuples(tuples, names=['first', 'second'])
+        >>> df = pd.DataFrame(np.random.randn(8, 2), index=idx, columns=['A', 'B'])
+        >>> df
+                             A         B
+        first second
+        bar   1.0     1.764052  0.400157
+              2.0     0.978738  2.240893
+        baz   1.0     1.867558 -0.977278
+              2.0     0.950088 -0.151357
+        foo   1.0    -0.103219  0.410599
+              2.0     0.144044  1.454274
+        qux   1.0     0.761038  0.121675
+              2.0     0.443863  0.333674
+        >>> outcome=df.change_index_dtype(dtype=str)
+        >>> outcome
+                             A         B
+        first second
+        bar   1.0     1.764052  0.400157
+              2.0     0.978738  2.240893
+        baz   1.0     1.867558 -0.977278
+              2.0     0.950088 -0.151357
+        foo   1.0    -0.103219  0.410599
+              2.0     0.144044  1.454274
+        qux   1.0     0.761038  0.121675
+              2.0     0.443863  0.333674
+        >>> outcome.index.dtypes
+        first     object
+        second    object
+        dtype: object
+        >>> outcome=df.change_index_dtype(dtype={'second':int})
+        >>> outcome
+                             A         B
+        first second
+        bar   1       1.764052  0.400157
+              2       0.978738  2.240893
+        baz   1       1.867558 -0.977278
+              2       0.950088 -0.151357
+        foo   1      -0.103219  0.410599
+              2       0.144044  1.454274
+        qux   1       0.761038  0.121675
+              2       0.443863  0.333674
+        >>> outcome.index.dtypes
+        first     object
+        second     int64
+        dtype: object
+        >>> outcome=df.change_index_dtype(dtype={0:'category',1:int})
+        >>> outcome
+                             A         B
+        first second
+        bar   1       1.764052  0.400157
+              2       0.978738  2.240893
+        baz   1       1.867558 -0.977278
+              2       0.950088 -0.151357
+        foo   1      -0.103219  0.410599
+              2       0.144044  1.454274
+        qux   1       0.761038  0.121675
+              2       0.443863  0.333674
+        >>> outcome.index.dtypes
+        first     category
+        second       int64
+        dtype: object
+
+    Args:
+        df: A pandas DataFrame.
+        dtype : Use a str or dtype to cast the entire Index
+            to the same type.
+            Alternatively, use a dictionary to change the MultiIndex
+            to new dtypes.
+        axis: Determines which axis to change the dtype(s).
+            Should be either 'index' or 'columns'.
+
+    Returns:
+        A pandas DataFrame with new Index.
+    """  # noqa: E501
+
+    check("axis", axis, [str])
+    if axis not in {"index", "columns"}:
+        raise ValueError("axis should be either index or columns.")
+
+    df = df[:]
+    current_index = getattr(df, axis)
+    if not isinstance(current_index, pd.MultiIndex):
+        if isinstance(dtype, dict):
+            raise TypeError(
+                "Changing the dtype via a dictionary "
+                "is not supported for a single index."
+            )
+        current_index = current_index.astype(dtype)
+        setattr(df, axis, current_index)
+        return df
+
+    if not isinstance(dtype, dict):
+        dtype = {
+            level_number: dtype
+            for level_number in range(current_index.nlevels)
+        }
+
+    all_str = all(isinstance(level, str) for level in dtype)
+    all_int = all(isinstance(level, int) for level in dtype)
+    if not all_str | all_int:
+        raise TypeError(
+            "The levels in the dictionary "
+            "should be either all strings or all integers."
+        )
+
+    dtype = {
+        current_index._get_level_number(label): _dtype
+        for label, _dtype in dtype.items()
+    }
+
+    new_levels = []
+    codes = current_index.codes
+    levels = current_index.levels
+
+    for level_number in range(current_index.nlevels):
+        _index = levels[level_number]
+        if level_number in dtype:
+            _dtype = dtype[level_number]
+            _index = _index.astype(_dtype)
+        new_levels.append(_index)
+
+    current_index = pd.MultiIndex(
+        levels=new_levels,
+        codes=codes,
+        names=current_index.names,
+        copy=False,
+        verify_integrity=False,
+    )
+    setattr(df, axis, current_index)
+    return df


Test patch:
diff --git a/tests/functions/test_change_index_dtype.py b/tests/functions/test_change_index_dtype.py
new file mode 100644
index 0000000..e261156
--- /dev/null
+++ b/tests/functions/test_change_index_dtype.py
@@ -0,0 +1,116 @@
+import numpy as np
+import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal
+
+
+@pytest.fixture
+def df():
+    """fixture dataframe"""
+    return pd.DataFrame(
+        {
+            "fam_id": [1, 1, 1, 2, 2, 2, 3, 3, 3],
+        }
+    )
+
+
+@pytest.fixture
+def df_multi():
+    """MultiIndex dataframe fixture."""
+    # https://stackoverflow.com/a/58328741/7175713
+    tuples = list(
+        zip(
+            *[
+                ["bar", "bar", "baz", "baz", "foo", "foo", "qux", "qux"],
+                [1.0, 2.0, 1.0, 2.0, 1.0, 2.0, 1.0, 2.0],
+            ]
+        )
+    )
+
+    idx = pd.MultiIndex.from_tuples(tuples, names=["first", "second"])
+    return pd.DataFrame(np.random.randn(8, 2), index=idx, columns=["A", "B"])
+
+
+def test_type_axis(df):
+    """Raise TypeError if wrong type is provided for axis."""
+    with pytest.raises(TypeError, match="axis should be one of.+"):
+        df.change_index_dtype(dtype=float, axis=1)
+
+
+def test_axis_values(df):
+    """Raise ValueError if wrong value is provided for axis."""
+    msg = "axis should be either index or columns."
+    with pytest.raises(ValueError, match=msg):
+        df.change_index_dtype(dtype=float, axis="INDEX")
+
+
+def test_type_dict_single_index(df):
+    """Raise TypeError if dtype is a dict and single index."""
+    msg = "Changing the dtype via a dictionary "
+    msg = "is not supported for a single index."
+    with pytest.raises(TypeError, match=msg):
+        df.change_index_dtype(dtype={0: float})
+
+
+def test_multiindex_diff_key_types(df_multi):
+    """Test output if the keys in the dictionary are not the same type."""
+    msg = "The levels in the dictionary "
+    msg = "should be either all strings or all integers."
+    with pytest.raises(TypeError, match=msg):
+        df_multi.change_index_dtype(dtype={1: int, "first": "category"})
+
+
+def test_multiindex_single_dtype(df_multi):
+    """Test output if MultiIndex and a single dtype is passed."""
+    actual = df_multi.change_index_dtype(dtype=str)
+    expected = (
+        df_multi.reset_index()
+        .astype({"second": str, "first": str})
+        .set_index(["first", "second"])
+    )
+    assert_frame_equal(expected, actual)
+
+
+def test_multiindex_single_key(df_multi):
+    """Test output if a dictionary is passed"""
+    actual = df_multi.change_index_dtype(dtype={"second": int})
+    expected = (
+        df_multi.reset_index("second")
+        .astype({"second": int})
+        .set_index("second", append=True)
+    )
+    assert_frame_equal(expected, actual)
+
+
+def test_multiindex_multiple_keys(df_multi):
+    """Test output if a dictionary is passed"""
+    actual = df_multi.change_index_dtype(
+        dtype={"second": int, "first": "category"}
+    )
+    expected = (
+        df_multi.reset_index()
+        .astype({"second": int, "first": "category"})
+        .set_index(["first", "second"])
+    )
+    assert_frame_equal(expected, actual)
+
+
+def test_multiindex_multiple_keys_columns(df_multi):
+    """Test output if a dictionary is passed"""
+    actual = df_multi.T.change_index_dtype(
+        dtype={"second": int, "first": "category"}, axis="columns"
+    )
+    expected = (
+        df_multi.reset_index()
+        .astype({"second": int, "first": "category"})
+        .set_index(["first", "second"])
+        .T
+    )
+    assert_frame_equal(expected, actual)
+
+
+def test_single_index(df):
+    """Test output if df.columns is not a multiindex"""
+    actual = df.change_index_dtype(dtype=float)
+    expected = df.set_index(np.arange(0.0, len(df)))
+    assert_frame_equal(expected, actual)