test_keys.py 2.2 KB

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