_hypothesis.py 2.3 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889
  1. """
  2. Hypothesis data generator helpers.
  3. """
  4. from datetime import datetime
  5. from hypothesis import strategies as st
  6. from hypothesis.extra.dateutil import timezones as dateutil_timezones
  7. from hypothesis.extra.pytz import timezones as pytz_timezones
  8. from pandas.compat import is_platform_windows
  9. import pandas as pd
  10. from pandas.tseries.offsets import (
  11. BMonthBegin,
  12. BMonthEnd,
  13. BQuarterBegin,
  14. BQuarterEnd,
  15. BYearBegin,
  16. BYearEnd,
  17. MonthBegin,
  18. MonthEnd,
  19. QuarterBegin,
  20. QuarterEnd,
  21. YearBegin,
  22. YearEnd,
  23. )
  24. OPTIONAL_INTS = st.lists(st.one_of(st.integers(), st.none()), max_size=10, min_size=3)
  25. OPTIONAL_FLOATS = st.lists(st.one_of(st.floats(), st.none()), max_size=10, min_size=3)
  26. OPTIONAL_TEXT = st.lists(st.one_of(st.none(), st.text()), max_size=10, min_size=3)
  27. OPTIONAL_DICTS = st.lists(
  28. st.one_of(st.none(), st.dictionaries(st.text(), st.integers())),
  29. max_size=10,
  30. min_size=3,
  31. )
  32. OPTIONAL_LISTS = st.lists(
  33. st.one_of(st.none(), st.lists(st.text(), max_size=10, min_size=3)),
  34. max_size=10,
  35. min_size=3,
  36. )
  37. OPTIONAL_ONE_OF_ALL = st.one_of(
  38. OPTIONAL_DICTS, OPTIONAL_FLOATS, OPTIONAL_INTS, OPTIONAL_LISTS, OPTIONAL_TEXT
  39. )
  40. if is_platform_windows():
  41. DATETIME_NO_TZ = st.datetimes(min_value=datetime(1900, 1, 1))
  42. else:
  43. DATETIME_NO_TZ = st.datetimes()
  44. DATETIME_JAN_1_1900_OPTIONAL_TZ = st.datetimes(
  45. min_value=pd.Timestamp(1900, 1, 1).to_pydatetime(),
  46. max_value=pd.Timestamp(1900, 1, 1).to_pydatetime(),
  47. timezones=st.one_of(st.none(), dateutil_timezones(), pytz_timezones()),
  48. )
  49. DATETIME_IN_PD_TIMESTAMP_RANGE_NO_TZ = st.datetimes(
  50. min_value=pd.Timestamp.min.to_pydatetime(warn=False),
  51. max_value=pd.Timestamp.max.to_pydatetime(warn=False),
  52. )
  53. INT_NEG_999_TO_POS_999 = st.integers(-999, 999)
  54. # The strategy for each type is registered in conftest.py, as they don't carry
  55. # enough runtime information (e.g. type hints) to infer how to build them.
  56. YQM_OFFSET = st.one_of(
  57. *map(
  58. st.from_type,
  59. [
  60. MonthBegin,
  61. MonthEnd,
  62. BMonthBegin,
  63. BMonthEnd,
  64. QuarterBegin,
  65. QuarterEnd,
  66. BQuarterBegin,
  67. BQuarterEnd,
  68. YearBegin,
  69. YearEnd,
  70. BYearBegin,
  71. BYearEnd,
  72. ],
  73. )
  74. )