test_keys.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. DataFrame,
  5. HDFStore,
  6. Index,
  7. Series,
  8. date_range,
  9. )
  10. from pandas.tests.io.pytables.common import (
  11. ensure_clean_store,
  12. tables,
  13. )
  14. pytestmark = [pytest.mark.single_cpu]
  15. def test_keys(setup_path):
  16. with ensure_clean_store(setup_path) as store:
  17. store["a"] = Series(
  18. np.arange(10, dtype=np.float64), index=date_range("2020-01-01", periods=10)
  19. )
  20. store["b"] = Series(
  21. range(10), dtype="float64", index=[f"i_{i}" for i in range(10)]
  22. )
  23. store["c"] = DataFrame(
  24. 1.1 * np.arange(120).reshape((30, 4)),
  25. columns=Index(list("ABCD"), dtype=object),
  26. index=Index([f"i-{i}" for i in range(30)], dtype=object),
  27. )
  28. assert len(store) == 3
  29. expected = {"/a", "/b", "/c"}
  30. assert set(store.keys()) == expected
  31. assert set(store) == expected
  32. def test_non_pandas_keys(tmp_path, setup_path):
  33. class Table1(tables.IsDescription):
  34. value1 = tables.Float32Col()
  35. class Table2(tables.IsDescription):
  36. value2 = tables.Float32Col()
  37. class Table3(tables.IsDescription):
  38. value3 = tables.Float32Col()
  39. path = tmp_path / setup_path
  40. with tables.open_file(path, mode="w") as h5file:
  41. group = h5file.create_group("/", "group")
  42. h5file.create_table(group, "table1", Table1, "Table 1")
  43. h5file.create_table(group, "table2", Table2, "Table 2")
  44. h5file.create_table(group, "table3", Table3, "Table 3")
  45. with HDFStore(path) as store:
  46. assert len(store.keys(include="native")) == 3
  47. expected = {"/group/table1", "/group/table2", "/group/table3"}
  48. assert set(store.keys(include="native")) == expected
  49. assert set(store.keys(include="pandas")) == set()
  50. for name in expected:
  51. df = store.get(name)
  52. assert len(df.columns) == 1
  53. def test_keys_illegal_include_keyword_value(setup_path):
  54. with ensure_clean_store(setup_path) as store:
  55. with pytest.raises(
  56. ValueError,
  57. match="`include` should be either 'pandas' or 'native' but is 'illegal'",
  58. ):
  59. store.keys(include="illegal")
  60. def test_keys_ignore_hdf_softlink(setup_path):
  61. # GH 20523
  62. # Puts a softlink into HDF file and rereads
  63. with ensure_clean_store(setup_path) as store:
  64. df = DataFrame({"A": range(5), "B": range(5)})
  65. store.put("df", df)
  66. assert store.keys() == ["/df"]
  67. store._handle.create_soft_link(store._handle.root, "symlink", "df")
  68. # Should ignore the softlink
  69. assert store.keys() == ["/df"]