1234567891011121314151617181920212223242526272829303132333435363738 |
- """
- Tests involving custom Index subclasses
- """
- import numpy as np
- from pandas import (
- DataFrame,
- Index,
- )
- import pandas._testing as tm
- class CustomIndex(Index):
- def __new__(cls, data, name=None):
- # assert that this index class cannot hold strings
- if any(isinstance(val, str) for val in data):
- raise TypeError("CustomIndex cannot hold strings")
- if name is None and hasattr(data, "name"):
- name = data.name
- data = np.array(data, dtype="O")
- return cls._simple_new(data, name)
- def test_insert_fallback_to_base_index():
- # https://github.com/pandas-dev/pandas/issues/47071
- idx = CustomIndex([1, 2, 3])
- result = idx.insert(0, "string")
- expected = Index(["string", 1, 2, 3], dtype=object)
- tm.assert_index_equal(result, expected)
- df = DataFrame(
- np.random.randn(2, 3), columns=idx, index=Index([1, 2], name="string")
- )
- result = df.reset_index()
- tm.assert_index_equal(result.columns, expected)
|