conftest.py 1.2 KB

12345678910111213141516171819202122232425262728293031
  1. import numpy as np
  2. import pytest
  3. from pandas import DataFrame, Index, MultiIndex
  4. from pandas.util import testing as tm
  5. @pytest.fixture
  6. def multiindex_dataframe_random_data():
  7. """DataFrame with 2 level MultiIndex with random data"""
  8. index = MultiIndex(levels=[['foo', 'bar', 'baz', 'qux'], ['one', 'two',
  9. 'three']],
  10. codes=[[0, 0, 0, 1, 1, 2, 2, 3, 3, 3],
  11. [0, 1, 2, 0, 1, 1, 2, 0, 1, 2]],
  12. names=['first', 'second'])
  13. return DataFrame(np.random.randn(10, 3), index=index,
  14. columns=Index(['A', 'B', 'C'], name='exp'))
  15. @pytest.fixture
  16. def multiindex_year_month_day_dataframe_random_data():
  17. """DataFrame with 3 level MultiIndex (year, month, day) covering
  18. first 100 business days from 2000-01-01 with random data"""
  19. tdf = tm.makeTimeDataFrame(100)
  20. ymd = tdf.groupby([lambda x: x.year, lambda x: x.month,
  21. lambda x: x.day]).sum()
  22. # use Int64Index, to make sure things work
  23. ymd.index.set_levels([lev.astype('i8') for lev in ymd.index.levels],
  24. inplace=True)
  25. ymd.index.set_names(['year', 'month', 'day'], inplace=True)
  26. return ymd