|
9 | 9 |
|
10 | 10 | import xarray as xr |
11 | 11 | from xarray.core.types import T_DuckArray |
12 | | -from xarray.core.utils import attempt_import |
| 12 | +from xarray.core.utils import attempt_import, module_available |
13 | 13 |
|
14 | 14 | if TYPE_CHECKING: |
15 | 15 | from xarray.core.types import _DTypeLikeNested, _ShapeLike |
|
22 | 22 |
|
23 | 23 | __all__ = [ |
24 | 24 | "attrs", |
| 25 | + "cftime_datetimes", |
| 26 | + "datetimes", |
25 | 27 | "dimension_names", |
26 | 28 | "dimension_sizes", |
27 | 29 | "names", |
@@ -84,6 +86,25 @@ def pandas_index_dtypes() -> st.SearchStrategy[np.dtype]: |
84 | 86 | ) |
85 | 87 |
|
86 | 88 |
|
| 89 | +def datetimes() -> st.SearchStrategy: |
| 90 | + """ |
| 91 | + Generates datetime objects including both standard library datetimes and cftime datetimes. |
| 92 | +
|
| 93 | + Returns standard library datetime.datetime objects, and if cftime is available, |
| 94 | + also includes cftime datetime objects from various calendars. |
| 95 | +
|
| 96 | + Requires the hypothesis package to be installed. |
| 97 | +
|
| 98 | + See Also |
| 99 | + -------- |
| 100 | + :ref:`testing.hypothesis`_ |
| 101 | + """ |
| 102 | + strategy = st.datetimes() |
| 103 | + if module_available("cftime"): |
| 104 | + strategy = strategy | cftime_datetimes() |
| 105 | + return strategy |
| 106 | + |
| 107 | + |
87 | 108 | # TODO Generalize to all valid unicode characters once formatting bugs in xarray's reprs are fixed + docs can handle it. |
88 | 109 | _readable_characters = st.characters( |
89 | 110 | categories=["L", "N"], max_codepoint=0x017F |
@@ -477,36 +498,41 @@ def unique_subset_of( |
477 | 498 | ) |
478 | 499 |
|
479 | 500 |
|
480 | | -class CFTimeStrategy(st.SearchStrategy): |
481 | | - def __init__(self, min_value, max_value): |
482 | | - super().__init__() |
483 | | - self.min_value = min_value |
484 | | - self.max_value = max_value |
| 501 | +@st.composite |
| 502 | +def cftime_datetimes(draw: st.DrawFn): |
| 503 | + """ |
| 504 | + Generates cftime datetime objects across various calendars. |
| 505 | +
|
| 506 | + This strategy generates cftime datetime objects from all available |
| 507 | + cftime calendars with dates ranging from year -99999 to 99999. |
| 508 | +
|
| 509 | + Requires both the hypothesis and cftime packages to be installed. |
| 510 | +
|
| 511 | + Returns |
| 512 | + ------- |
| 513 | + cftime_datetime_strategy |
| 514 | + Strategy for generating cftime datetime objects. |
| 515 | +
|
| 516 | + See Also |
| 517 | + -------- |
| 518 | + :ref:`testing.hypothesis`_ |
| 519 | + """ |
| 520 | + from xarray.tests import _all_cftime_date_types |
| 521 | + |
| 522 | + date_types = _all_cftime_date_types() |
| 523 | + calendars = list(date_types) |
| 524 | + |
| 525 | + calendar = draw(st.sampled_from(calendars)) |
| 526 | + date_type = date_types[calendar] |
| 527 | + |
| 528 | + with warnings.catch_warnings(): |
| 529 | + warnings.filterwarnings("ignore", message=".*date/calendar/year zero.*") |
| 530 | + daysinmonth = date_type(99999, 12, 1).daysinmonth |
| 531 | + min_value = date_type(-99999, 1, 1) |
| 532 | + max_value = date_type(99999, 12, daysinmonth, 23, 59, 59, 999999) |
485 | 533 |
|
486 | | - def do_draw(self, data): |
487 | 534 | unit_microsecond = datetime.timedelta(microseconds=1) |
488 | | - timespan_microseconds = (self.max_value - self.min_value) // unit_microsecond |
489 | | - result = data.draw_integer(0, timespan_microseconds) |
490 | | - with warnings.catch_warnings(): |
491 | | - warnings.filterwarnings("ignore", message=".*date/calendar/year zero.*") |
492 | | - return self.min_value + datetime.timedelta(microseconds=result) |
493 | | - |
494 | | - |
495 | | -class CFTimeStrategyISO8601(st.SearchStrategy): |
496 | | - def __init__(self): |
497 | | - from xarray.tests.test_coding_times import _all_cftime_date_types |
498 | | - |
499 | | - super().__init__() |
500 | | - self.date_types = _all_cftime_date_types() |
501 | | - self.calendars = list(self.date_types) |
502 | | - |
503 | | - def do_draw(self, data): |
504 | | - calendar = data.draw(st.sampled_from(self.calendars)) |
505 | | - date_type = self.date_types[calendar] |
506 | | - with warnings.catch_warnings(): |
507 | | - warnings.filterwarnings("ignore", message=".*date/calendar/year zero.*") |
508 | | - daysinmonth = date_type(99999, 12, 1).daysinmonth |
509 | | - min_value = date_type(-99999, 1, 1) |
510 | | - max_value = date_type(99999, 12, daysinmonth, 23, 59, 59, 999999) |
511 | | - strategy = CFTimeStrategy(min_value, max_value) |
512 | | - return strategy.do_draw(data) |
| 535 | + timespan_microseconds = (max_value - min_value) // unit_microsecond |
| 536 | + microseconds_offset = draw(st.integers(0, timespan_microseconds)) |
| 537 | + |
| 538 | + return min_value + datetime.timedelta(microseconds=microseconds_offset) |
0 commit comments