File size: 1,058 Bytes
7885a28 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 |
"""
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.default_rng(2).standard_normal((2, 3)),
columns=idx,
index=Index([1, 2], name="string"),
)
result = df.reset_index()
tm.assert_index_equal(result.columns, expected)
|