conftest.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221
  1. import numpy as np
  2. import pytest
  3. from pandas import DataFrame, NaT, compat, date_range
  4. import pandas.util.testing as tm
  5. @pytest.fixture
  6. def float_frame():
  7. """
  8. Fixture for DataFrame of floats with index of unique strings
  9. Columns are ['A', 'B', 'C', 'D'].
  10. """
  11. return DataFrame(tm.getSeriesData())
  12. @pytest.fixture
  13. def float_frame_with_na():
  14. """
  15. Fixture for DataFrame of floats with index of unique strings
  16. Columns are ['A', 'B', 'C', 'D']; some entries are missing
  17. """
  18. df = DataFrame(tm.getSeriesData())
  19. # set some NAs
  20. df.loc[5:10] = np.nan
  21. df.loc[15:20, -2:] = np.nan
  22. return df
  23. @pytest.fixture
  24. def float_frame2():
  25. """
  26. Fixture for DataFrame of floats with index of unique strings
  27. Columns are ['D', 'C', 'B', 'A']
  28. """
  29. return DataFrame(tm.getSeriesData(), columns=['D', 'C', 'B', 'A'])
  30. @pytest.fixture
  31. def bool_frame_with_na():
  32. """
  33. Fixture for DataFrame of booleans with index of unique strings
  34. Columns are ['A', 'B', 'C', 'D']; some entries are missing
  35. """
  36. df = DataFrame(tm.getSeriesData()) > 0
  37. df = df.astype(object)
  38. # set some NAs
  39. df.loc[5:10] = np.nan
  40. df.loc[15:20, -2:] = np.nan
  41. return df
  42. @pytest.fixture
  43. def int_frame():
  44. """
  45. Fixture for DataFrame of ints with index of unique strings
  46. Columns are ['A', 'B', 'C', 'D']
  47. """
  48. df = DataFrame({k: v.astype(int)
  49. for k, v in compat.iteritems(tm.getSeriesData())})
  50. # force these all to int64 to avoid platform testing issues
  51. return DataFrame({c: s for c, s in compat.iteritems(df)}, dtype=np.int64)
  52. @pytest.fixture
  53. def datetime_frame():
  54. """
  55. Fixture for DataFrame of floats with DatetimeIndex
  56. Columns are ['A', 'B', 'C', 'D']
  57. """
  58. return DataFrame(tm.getTimeSeriesData())
  59. @pytest.fixture
  60. def float_string_frame():
  61. """
  62. Fixture for DataFrame of floats and strings with index of unique strings
  63. Columns are ['A', 'B', 'C', 'D', 'foo'].
  64. """
  65. df = DataFrame(tm.getSeriesData())
  66. df['foo'] = 'bar'
  67. return df
  68. @pytest.fixture
  69. def mixed_float_frame():
  70. """
  71. Fixture for DataFrame of different float types with index of unique strings
  72. Columns are ['A', 'B', 'C', 'D'].
  73. """
  74. df = DataFrame(tm.getSeriesData())
  75. df.A = df.A.astype('float32')
  76. df.B = df.B.astype('float32')
  77. df.C = df.C.astype('float16')
  78. df.D = df.D.astype('float64')
  79. return df
  80. @pytest.fixture
  81. def mixed_float_frame2():
  82. """
  83. Fixture for DataFrame of different float types with index of unique strings
  84. Columns are ['A', 'B', 'C', 'D'].
  85. """
  86. df = DataFrame(tm.getSeriesData())
  87. df.D = df.D.astype('float32')
  88. df.C = df.C.astype('float32')
  89. df.B = df.B.astype('float16')
  90. df.D = df.D.astype('float64')
  91. return df
  92. @pytest.fixture
  93. def mixed_int_frame():
  94. """
  95. Fixture for DataFrame of different int types with index of unique strings
  96. Columns are ['A', 'B', 'C', 'D'].
  97. """
  98. df = DataFrame({k: v.astype(int)
  99. for k, v in compat.iteritems(tm.getSeriesData())})
  100. df.A = df.A.astype('int32')
  101. df.B = np.ones(len(df.B), dtype='uint64')
  102. df.C = df.C.astype('uint8')
  103. df.D = df.C.astype('int64')
  104. return df
  105. @pytest.fixture
  106. def mixed_type_frame():
  107. """
  108. Fixture for DataFrame of float/int/string columns with RangeIndex
  109. Columns are ['a', 'b', 'c', 'float32', 'int32'].
  110. """
  111. return DataFrame({'a': 1., 'b': 2, 'c': 'foo',
  112. 'float32': np.array([1.] * 10, dtype='float32'),
  113. 'int32': np.array([1] * 10, dtype='int32')},
  114. index=np.arange(10))
  115. @pytest.fixture
  116. def timezone_frame():
  117. """
  118. Fixture for DataFrame of date_range Series with different time zones
  119. Columns are ['A', 'B', 'C']; some entries are missing
  120. """
  121. df = DataFrame({'A': date_range('20130101', periods=3),
  122. 'B': date_range('20130101', periods=3,
  123. tz='US/Eastern'),
  124. 'C': date_range('20130101', periods=3,
  125. tz='CET')})
  126. df.iloc[1, 1] = NaT
  127. df.iloc[1, 2] = NaT
  128. return df
  129. @pytest.fixture
  130. def empty_frame():
  131. """
  132. Fixture for empty DataFrame
  133. """
  134. return DataFrame({})
  135. @pytest.fixture
  136. def datetime_series():
  137. """
  138. Fixture for Series of floats with DatetimeIndex
  139. """
  140. return tm.makeTimeSeries(nper=30)
  141. @pytest.fixture
  142. def datetime_series_short():
  143. """
  144. Fixture for Series of floats with DatetimeIndex
  145. """
  146. return tm.makeTimeSeries(nper=30)[5:]
  147. @pytest.fixture
  148. def simple_frame():
  149. """
  150. Fixture for simple 3x3 DataFrame
  151. Columns are ['one', 'two', 'three'], index is ['a', 'b', 'c'].
  152. """
  153. arr = np.array([[1., 2., 3.],
  154. [4., 5., 6.],
  155. [7., 8., 9.]])
  156. return DataFrame(arr, columns=['one', 'two', 'three'],
  157. index=['a', 'b', 'c'])
  158. @pytest.fixture
  159. def frame_of_index_cols():
  160. """
  161. Fixture for DataFrame of columns that can be used for indexing
  162. Columns are ['A', 'B', 'C', 'D', 'E', ('tuple', 'as', 'label')];
  163. 'A' & 'B' contain duplicates (but are jointly unique), the rest are unique.
  164. """
  165. df = DataFrame({'A': ['foo', 'foo', 'foo', 'bar', 'bar'],
  166. 'B': ['one', 'two', 'three', 'one', 'two'],
  167. 'C': ['a', 'b', 'c', 'd', 'e'],
  168. 'D': np.random.randn(5),
  169. 'E': np.random.randn(5),
  170. ('tuple', 'as', 'label'): np.random.randn(5)})
  171. return df