conftest.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100
  1. import numpy as np
  2. import pytest
  3. from pandas import (
  4. DataFrame,
  5. Index,
  6. NaT,
  7. date_range,
  8. )
  9. @pytest.fixture
  10. def datetime_frame() -> DataFrame:
  11. """
  12. Fixture for DataFrame of floats with DatetimeIndex
  13. Columns are ['A', 'B', 'C', 'D']
  14. """
  15. return DataFrame(
  16. np.random.default_rng(2).standard_normal((100, 4)),
  17. columns=Index(list("ABCD")),
  18. index=date_range("2000-01-01", periods=100, freq="B"),
  19. )
  20. @pytest.fixture
  21. def float_string_frame():
  22. """
  23. Fixture for DataFrame of floats and strings with index of unique strings
  24. Columns are ['A', 'B', 'C', 'D', 'foo'].
  25. """
  26. df = DataFrame(
  27. np.random.default_rng(2).standard_normal((30, 4)),
  28. index=Index([f"foo_{i}" for i in range(30)], dtype=object),
  29. columns=Index(list("ABCD")),
  30. )
  31. df["foo"] = "bar"
  32. return df
  33. @pytest.fixture
  34. def mixed_float_frame():
  35. """
  36. Fixture for DataFrame of different float types with index of unique strings
  37. Columns are ['A', 'B', 'C', 'D'].
  38. """
  39. df = DataFrame(
  40. {
  41. col: np.random.default_rng(2).random(30, dtype=dtype)
  42. for col, dtype in zip(
  43. list("ABCD"), ["float32", "float32", "float32", "float64"]
  44. )
  45. },
  46. index=Index([f"foo_{i}" for i in range(30)], dtype=object),
  47. )
  48. # not supported by numpy random
  49. df["C"] = df["C"].astype("float16")
  50. return df
  51. @pytest.fixture
  52. def mixed_int_frame():
  53. """
  54. Fixture for DataFrame of different int types with index of unique strings
  55. Columns are ['A', 'B', 'C', 'D'].
  56. """
  57. return DataFrame(
  58. {
  59. col: np.ones(30, dtype=dtype)
  60. for col, dtype in zip(list("ABCD"), ["int32", "uint64", "uint8", "int64"])
  61. },
  62. index=Index([f"foo_{i}" for i in range(30)], dtype=object),
  63. )
  64. @pytest.fixture
  65. def timezone_frame():
  66. """
  67. Fixture for DataFrame of date_range Series with different time zones
  68. Columns are ['A', 'B', 'C']; some entries are missing
  69. A B C
  70. 0 2013-01-01 2013-01-01 00:00:00-05:00 2013-01-01 00:00:00+01:00
  71. 1 2013-01-02 NaT NaT
  72. 2 2013-01-03 2013-01-03 00:00:00-05:00 2013-01-03 00:00:00+01:00
  73. """
  74. df = DataFrame(
  75. {
  76. "A": date_range("20130101", periods=3),
  77. "B": date_range("20130101", periods=3, tz="US/Eastern"),
  78. "C": date_range("20130101", periods=3, tz="CET"),
  79. }
  80. )
  81. df.iloc[1, 1] = NaT
  82. df.iloc[1, 2] = NaT
  83. return df