test_get_dummies.py 1.6 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253
  1. import numpy as np
  2. from pandas import (
  3. DataFrame,
  4. Index,
  5. MultiIndex,
  6. Series,
  7. _testing as tm,
  8. )
  9. def test_get_dummies(any_string_dtype):
  10. s = Series(["a|b", "a|c", np.nan], dtype=any_string_dtype)
  11. result = s.str.get_dummies("|")
  12. expected = DataFrame([[1, 1, 0], [1, 0, 1], [0, 0, 0]], columns=list("abc"))
  13. tm.assert_frame_equal(result, expected)
  14. s = Series(["a;b", "a", 7], dtype=any_string_dtype)
  15. result = s.str.get_dummies(";")
  16. expected = DataFrame([[0, 1, 1], [0, 1, 0], [1, 0, 0]], columns=list("7ab"))
  17. tm.assert_frame_equal(result, expected)
  18. def test_get_dummies_index():
  19. # GH9980, GH8028
  20. idx = Index(["a|b", "a|c", "b|c"])
  21. result = idx.str.get_dummies("|")
  22. expected = MultiIndex.from_tuples(
  23. [(1, 1, 0), (1, 0, 1), (0, 1, 1)], names=("a", "b", "c")
  24. )
  25. tm.assert_index_equal(result, expected)
  26. def test_get_dummies_with_name_dummy(any_string_dtype):
  27. # GH 12180
  28. # Dummies named 'name' should work as expected
  29. s = Series(["a", "b,name", "b"], dtype=any_string_dtype)
  30. result = s.str.get_dummies(",")
  31. expected = DataFrame([[1, 0, 0], [0, 1, 1], [0, 1, 0]], columns=["a", "b", "name"])
  32. tm.assert_frame_equal(result, expected)
  33. def test_get_dummies_with_name_dummy_index():
  34. # GH 12180
  35. # Dummies named 'name' should work as expected
  36. idx = Index(["a|b", "name|c", "b|name"])
  37. result = idx.str.get_dummies("|")
  38. expected = MultiIndex.from_tuples(
  39. [(1, 1, 0, 0), (0, 0, 1, 1), (0, 1, 0, 1)], names=("a", "b", "c", "name")
  40. )
  41. tm.assert_index_equal(result, expected)