test_subclass.py 1014 B

1234567891011121314151617181920212223242526272829303132333435363738
  1. """
  2. Tests involving custom Index subclasses
  3. """
  4. import numpy as np
  5. from pandas import (
  6. DataFrame,
  7. Index,
  8. )
  9. import pandas._testing as tm
  10. class CustomIndex(Index):
  11. def __new__(cls, data, name=None):
  12. # assert that this index class cannot hold strings
  13. if any(isinstance(val, str) for val in data):
  14. raise TypeError("CustomIndex cannot hold strings")
  15. if name is None and hasattr(data, "name"):
  16. name = data.name
  17. data = np.array(data, dtype="O")
  18. return cls._simple_new(data, name)
  19. def test_insert_fallback_to_base_index():
  20. # https://github.com/pandas-dev/pandas/issues/47071
  21. idx = CustomIndex([1, 2, 3])
  22. result = idx.insert(0, "string")
  23. expected = Index(["string", 1, 2, 3], dtype=object)
  24. tm.assert_index_equal(result, expected)
  25. df = DataFrame(
  26. np.random.randn(2, 3), columns=idx, index=Index([1, 2], name="string")
  27. )
  28. result = df.reset_index()
  29. tm.assert_index_equal(result.columns, expected)