1from itertools import product
2
3import numpy as np
4import pandas as pd
5import pytest
6
7from xarray import CFTimeIndex
8from xarray.coding.cftime_offsets import (
9    _MONTH_ABBREVIATIONS,
10    BaseCFTimeOffset,
11    Day,
12    Hour,
13    Microsecond,
14    Millisecond,
15    Minute,
16    MonthBegin,
17    MonthEnd,
18    QuarterBegin,
19    QuarterEnd,
20    Second,
21    YearBegin,
22    YearEnd,
23    _days_in_month,
24    cftime_range,
25    get_date_type,
26    to_cftime_datetime,
27    to_offset,
28)
29from xarray.tests import _CFTIME_CALENDARS
30
31cftime = pytest.importorskip("cftime")
32
33
34def _id_func(param):
35    """Called on each parameter passed to pytest.mark.parametrize"""
36    return str(param)
37
38
39@pytest.fixture(params=_CFTIME_CALENDARS)
40def calendar(request):
41    return request.param
42
43
44@pytest.mark.parametrize(
45    ("offset", "expected_n"),
46    [
47        (BaseCFTimeOffset(), 1),
48        (YearBegin(), 1),
49        (YearEnd(), 1),
50        (QuarterBegin(), 1),
51        (QuarterEnd(), 1),
52        (BaseCFTimeOffset(n=2), 2),
53        (YearBegin(n=2), 2),
54        (YearEnd(n=2), 2),
55        (QuarterBegin(n=2), 2),
56        (QuarterEnd(n=2), 2),
57    ],
58    ids=_id_func,
59)
60def test_cftime_offset_constructor_valid_n(offset, expected_n):
61    assert offset.n == expected_n
62
63
64@pytest.mark.parametrize(
65    ("offset", "invalid_n"),
66    [
67        (BaseCFTimeOffset, 1.5),
68        (YearBegin, 1.5),
69        (YearEnd, 1.5),
70        (QuarterBegin, 1.5),
71        (QuarterEnd, 1.5),
72    ],
73    ids=_id_func,
74)
75def test_cftime_offset_constructor_invalid_n(offset, invalid_n):
76    with pytest.raises(TypeError):
77        offset(n=invalid_n)
78
79
80@pytest.mark.parametrize(
81    ("offset", "expected_month"),
82    [
83        (YearBegin(), 1),
84        (YearEnd(), 12),
85        (YearBegin(month=5), 5),
86        (YearEnd(month=5), 5),
87        (QuarterBegin(), 3),
88        (QuarterEnd(), 3),
89        (QuarterBegin(month=5), 5),
90        (QuarterEnd(month=5), 5),
91    ],
92    ids=_id_func,
93)
94def test_year_offset_constructor_valid_month(offset, expected_month):
95    assert offset.month == expected_month
96
97
98@pytest.mark.parametrize(
99    ("offset", "invalid_month", "exception"),
100    [
101        (YearBegin, 0, ValueError),
102        (YearEnd, 0, ValueError),
103        (YearBegin, 13, ValueError),
104        (YearEnd, 13, ValueError),
105        (YearBegin, 1.5, TypeError),
106        (YearEnd, 1.5, TypeError),
107        (QuarterBegin, 0, ValueError),
108        (QuarterEnd, 0, ValueError),
109        (QuarterBegin, 1.5, TypeError),
110        (QuarterEnd, 1.5, TypeError),
111        (QuarterBegin, 13, ValueError),
112        (QuarterEnd, 13, ValueError),
113    ],
114    ids=_id_func,
115)
116def test_year_offset_constructor_invalid_month(offset, invalid_month, exception):
117    with pytest.raises(exception):
118        offset(month=invalid_month)
119
120
121@pytest.mark.parametrize(
122    ("offset", "expected"),
123    [
124        (BaseCFTimeOffset(), None),
125        (MonthBegin(), "MS"),
126        (YearBegin(), "AS-JAN"),
127        (QuarterBegin(), "QS-MAR"),
128    ],
129    ids=_id_func,
130)
131def test_rule_code(offset, expected):
132    assert offset.rule_code() == expected
133
134
135@pytest.mark.parametrize(
136    ("offset", "expected"),
137    [
138        (BaseCFTimeOffset(), "<BaseCFTimeOffset: n=1>"),
139        (YearBegin(), "<YearBegin: n=1, month=1>"),
140        (QuarterBegin(), "<QuarterBegin: n=1, month=3>"),
141    ],
142    ids=_id_func,
143)
144def test_str_and_repr(offset, expected):
145    assert str(offset) == expected
146    assert repr(offset) == expected
147
148
149@pytest.mark.parametrize(
150    "offset",
151    [BaseCFTimeOffset(), MonthBegin(), QuarterBegin(), YearBegin()],
152    ids=_id_func,
153)
154def test_to_offset_offset_input(offset):
155    assert to_offset(offset) == offset
156
157
158@pytest.mark.parametrize(
159    ("freq", "expected"),
160    [
161        ("M", MonthEnd()),
162        ("2M", MonthEnd(n=2)),
163        ("MS", MonthBegin()),
164        ("2MS", MonthBegin(n=2)),
165        ("D", Day()),
166        ("2D", Day(n=2)),
167        ("H", Hour()),
168        ("2H", Hour(n=2)),
169        ("T", Minute()),
170        ("2T", Minute(n=2)),
171        ("min", Minute()),
172        ("2min", Minute(n=2)),
173        ("S", Second()),
174        ("2S", Second(n=2)),
175        ("L", Millisecond(n=1)),
176        ("2L", Millisecond(n=2)),
177        ("ms", Millisecond(n=1)),
178        ("2ms", Millisecond(n=2)),
179        ("U", Microsecond(n=1)),
180        ("2U", Microsecond(n=2)),
181        ("us", Microsecond(n=1)),
182        ("2us", Microsecond(n=2)),
183    ],
184    ids=_id_func,
185)
186def test_to_offset_sub_annual(freq, expected):
187    assert to_offset(freq) == expected
188
189
190_ANNUAL_OFFSET_TYPES = {"A": YearEnd, "AS": YearBegin}
191
192
193@pytest.mark.parametrize(
194    ("month_int", "month_label"), list(_MONTH_ABBREVIATIONS.items()) + [(0, "")]
195)
196@pytest.mark.parametrize("multiple", [None, 2])
197@pytest.mark.parametrize("offset_str", ["AS", "A"])
198def test_to_offset_annual(month_label, month_int, multiple, offset_str):
199    freq = offset_str
200    offset_type = _ANNUAL_OFFSET_TYPES[offset_str]
201    if month_label:
202        freq = "-".join([freq, month_label])
203    if multiple:
204        freq = f"{multiple}{freq}"
205    result = to_offset(freq)
206
207    if multiple and month_int:
208        expected = offset_type(n=multiple, month=month_int)
209    elif multiple:
210        expected = offset_type(n=multiple)
211    elif month_int:
212        expected = offset_type(month=month_int)
213    else:
214        expected = offset_type()
215    assert result == expected
216
217
218_QUARTER_OFFSET_TYPES = {"Q": QuarterEnd, "QS": QuarterBegin}
219
220
221@pytest.mark.parametrize(
222    ("month_int", "month_label"), list(_MONTH_ABBREVIATIONS.items()) + [(0, "")]
223)
224@pytest.mark.parametrize("multiple", [None, 2])
225@pytest.mark.parametrize("offset_str", ["QS", "Q"])
226def test_to_offset_quarter(month_label, month_int, multiple, offset_str):
227    freq = offset_str
228    offset_type = _QUARTER_OFFSET_TYPES[offset_str]
229    if month_label:
230        freq = "-".join([freq, month_label])
231    if multiple:
232        freq = f"{multiple}{freq}"
233    result = to_offset(freq)
234
235    if multiple and month_int:
236        expected = offset_type(n=multiple, month=month_int)
237    elif multiple:
238        if month_int:
239            expected = offset_type(n=multiple)
240        else:
241            if offset_type == QuarterBegin:
242                expected = offset_type(n=multiple, month=1)
243            elif offset_type == QuarterEnd:
244                expected = offset_type(n=multiple, month=12)
245    elif month_int:
246        expected = offset_type(month=month_int)
247    else:
248        if offset_type == QuarterBegin:
249            expected = offset_type(month=1)
250        elif offset_type == QuarterEnd:
251            expected = offset_type(month=12)
252    assert result == expected
253
254
255@pytest.mark.parametrize("freq", ["Z", "7min2", "AM", "M-", "AS-", "QS-", "1H1min"])
256def test_invalid_to_offset_str(freq):
257    with pytest.raises(ValueError):
258        to_offset(freq)
259
260
261@pytest.mark.parametrize(
262    ("argument", "expected_date_args"),
263    [("2000-01-01", (2000, 1, 1)), ((2000, 1, 1), (2000, 1, 1))],
264    ids=_id_func,
265)
266def test_to_cftime_datetime(calendar, argument, expected_date_args):
267    date_type = get_date_type(calendar)
268    expected = date_type(*expected_date_args)
269    if isinstance(argument, tuple):
270        argument = date_type(*argument)
271    result = to_cftime_datetime(argument, calendar=calendar)
272    assert result == expected
273
274
275def test_to_cftime_datetime_error_no_calendar():
276    with pytest.raises(ValueError):
277        to_cftime_datetime("2000")
278
279
280def test_to_cftime_datetime_error_type_error():
281    with pytest.raises(TypeError):
282        to_cftime_datetime(1)
283
284
285_EQ_TESTS_A = [
286    BaseCFTimeOffset(),
287    YearBegin(),
288    YearEnd(),
289    YearBegin(month=2),
290    YearEnd(month=2),
291    QuarterBegin(),
292    QuarterEnd(),
293    QuarterBegin(month=2),
294    QuarterEnd(month=2),
295    MonthBegin(),
296    MonthEnd(),
297    Day(),
298    Hour(),
299    Minute(),
300    Second(),
301    Millisecond(),
302    Microsecond(),
303]
304_EQ_TESTS_B = [
305    BaseCFTimeOffset(n=2),
306    YearBegin(n=2),
307    YearEnd(n=2),
308    YearBegin(n=2, month=2),
309    YearEnd(n=2, month=2),
310    QuarterBegin(n=2),
311    QuarterEnd(n=2),
312    QuarterBegin(n=2, month=2),
313    QuarterEnd(n=2, month=2),
314    MonthBegin(n=2),
315    MonthEnd(n=2),
316    Day(n=2),
317    Hour(n=2),
318    Minute(n=2),
319    Second(n=2),
320    Millisecond(n=2),
321    Microsecond(n=2),
322]
323
324
325@pytest.mark.parametrize(("a", "b"), product(_EQ_TESTS_A, _EQ_TESTS_B), ids=_id_func)
326def test_neq(a, b):
327    assert a != b
328
329
330_EQ_TESTS_B_COPY = [
331    BaseCFTimeOffset(n=2),
332    YearBegin(n=2),
333    YearEnd(n=2),
334    YearBegin(n=2, month=2),
335    YearEnd(n=2, month=2),
336    QuarterBegin(n=2),
337    QuarterEnd(n=2),
338    QuarterBegin(n=2, month=2),
339    QuarterEnd(n=2, month=2),
340    MonthBegin(n=2),
341    MonthEnd(n=2),
342    Day(n=2),
343    Hour(n=2),
344    Minute(n=2),
345    Second(n=2),
346    Millisecond(n=2),
347    Microsecond(n=2),
348]
349
350
351@pytest.mark.parametrize(("a", "b"), zip(_EQ_TESTS_B, _EQ_TESTS_B_COPY), ids=_id_func)
352def test_eq(a, b):
353    assert a == b
354
355
356_MUL_TESTS = [
357    (BaseCFTimeOffset(), BaseCFTimeOffset(n=3)),
358    (YearEnd(), YearEnd(n=3)),
359    (YearBegin(), YearBegin(n=3)),
360    (QuarterEnd(), QuarterEnd(n=3)),
361    (QuarterBegin(), QuarterBegin(n=3)),
362    (MonthEnd(), MonthEnd(n=3)),
363    (MonthBegin(), MonthBegin(n=3)),
364    (Day(), Day(n=3)),
365    (Hour(), Hour(n=3)),
366    (Minute(), Minute(n=3)),
367    (Second(), Second(n=3)),
368    (Millisecond(), Millisecond(n=3)),
369    (Microsecond(), Microsecond(n=3)),
370]
371
372
373@pytest.mark.parametrize(("offset", "expected"), _MUL_TESTS, ids=_id_func)
374def test_mul(offset, expected):
375    assert offset * 3 == expected
376
377
378@pytest.mark.parametrize(("offset", "expected"), _MUL_TESTS, ids=_id_func)
379def test_rmul(offset, expected):
380    assert 3 * offset == expected
381
382
383@pytest.mark.parametrize(
384    ("offset", "expected"),
385    [
386        (BaseCFTimeOffset(), BaseCFTimeOffset(n=-1)),
387        (YearEnd(), YearEnd(n=-1)),
388        (YearBegin(), YearBegin(n=-1)),
389        (QuarterEnd(), QuarterEnd(n=-1)),
390        (QuarterBegin(), QuarterBegin(n=-1)),
391        (MonthEnd(), MonthEnd(n=-1)),
392        (MonthBegin(), MonthBegin(n=-1)),
393        (Day(), Day(n=-1)),
394        (Hour(), Hour(n=-1)),
395        (Minute(), Minute(n=-1)),
396        (Second(), Second(n=-1)),
397        (Millisecond(), Millisecond(n=-1)),
398        (Microsecond(), Microsecond(n=-1)),
399    ],
400    ids=_id_func,
401)
402def test_neg(offset, expected):
403    assert -offset == expected
404
405
406_ADD_TESTS = [
407    (Day(n=2), (1, 1, 3)),
408    (Hour(n=2), (1, 1, 1, 2)),
409    (Minute(n=2), (1, 1, 1, 0, 2)),
410    (Second(n=2), (1, 1, 1, 0, 0, 2)),
411    (Millisecond(n=2), (1, 1, 1, 0, 0, 0, 2000)),
412    (Microsecond(n=2), (1, 1, 1, 0, 0, 0, 2)),
413]
414
415
416@pytest.mark.parametrize(("offset", "expected_date_args"), _ADD_TESTS, ids=_id_func)
417def test_add_sub_monthly(offset, expected_date_args, calendar):
418    date_type = get_date_type(calendar)
419    initial = date_type(1, 1, 1)
420    expected = date_type(*expected_date_args)
421    result = offset + initial
422    assert result == expected
423
424
425@pytest.mark.parametrize(("offset", "expected_date_args"), _ADD_TESTS, ids=_id_func)
426def test_radd_sub_monthly(offset, expected_date_args, calendar):
427    date_type = get_date_type(calendar)
428    initial = date_type(1, 1, 1)
429    expected = date_type(*expected_date_args)
430    result = initial + offset
431    assert result == expected
432
433
434@pytest.mark.parametrize(
435    ("offset", "expected_date_args"),
436    [
437        (Day(n=2), (1, 1, 1)),
438        (Hour(n=2), (1, 1, 2, 22)),
439        (Minute(n=2), (1, 1, 2, 23, 58)),
440        (Second(n=2), (1, 1, 2, 23, 59, 58)),
441        (Millisecond(n=2), (1, 1, 2, 23, 59, 59, 998000)),
442        (Microsecond(n=2), (1, 1, 2, 23, 59, 59, 999998)),
443    ],
444    ids=_id_func,
445)
446def test_rsub_sub_monthly(offset, expected_date_args, calendar):
447    date_type = get_date_type(calendar)
448    initial = date_type(1, 1, 3)
449    expected = date_type(*expected_date_args)
450    result = initial - offset
451    assert result == expected
452
453
454@pytest.mark.parametrize("offset", _EQ_TESTS_A, ids=_id_func)
455def test_sub_error(offset, calendar):
456    date_type = get_date_type(calendar)
457    initial = date_type(1, 1, 1)
458    with pytest.raises(TypeError):
459        offset - initial
460
461
462@pytest.mark.parametrize(("a", "b"), zip(_EQ_TESTS_A, _EQ_TESTS_B), ids=_id_func)
463def test_minus_offset(a, b):
464    result = b - a
465    expected = a
466    assert result == expected
467
468
469@pytest.mark.parametrize(
470    ("a", "b"),
471    list(zip(np.roll(_EQ_TESTS_A, 1), _EQ_TESTS_B))  # type: ignore[arg-type]
472    + [(YearEnd(month=1), YearEnd(month=2))],
473    ids=_id_func,
474)
475def test_minus_offset_error(a, b):
476    with pytest.raises(TypeError):
477        b - a
478
479
480def test_days_in_month_non_december(calendar):
481    date_type = get_date_type(calendar)
482    reference = date_type(1, 4, 1)
483    assert _days_in_month(reference) == 30
484
485
486def test_days_in_month_december(calendar):
487    if calendar == "360_day":
488        expected = 30
489    else:
490        expected = 31
491    date_type = get_date_type(calendar)
492    reference = date_type(1, 12, 5)
493    assert _days_in_month(reference) == expected
494
495
496@pytest.mark.parametrize(
497    ("initial_date_args", "offset", "expected_date_args"),
498    [
499        ((1, 1, 1), MonthBegin(), (1, 2, 1)),
500        ((1, 1, 1), MonthBegin(n=2), (1, 3, 1)),
501        ((1, 1, 7), MonthBegin(), (1, 2, 1)),
502        ((1, 1, 7), MonthBegin(n=2), (1, 3, 1)),
503        ((1, 3, 1), MonthBegin(n=-1), (1, 2, 1)),
504        ((1, 3, 1), MonthBegin(n=-2), (1, 1, 1)),
505        ((1, 3, 3), MonthBegin(n=-1), (1, 3, 1)),
506        ((1, 3, 3), MonthBegin(n=-2), (1, 2, 1)),
507        ((1, 2, 1), MonthBegin(n=14), (2, 4, 1)),
508        ((2, 4, 1), MonthBegin(n=-14), (1, 2, 1)),
509        ((1, 1, 1, 5, 5, 5, 5), MonthBegin(), (1, 2, 1, 5, 5, 5, 5)),
510        ((1, 1, 3, 5, 5, 5, 5), MonthBegin(), (1, 2, 1, 5, 5, 5, 5)),
511        ((1, 1, 3, 5, 5, 5, 5), MonthBegin(n=-1), (1, 1, 1, 5, 5, 5, 5)),
512    ],
513    ids=_id_func,
514)
515def test_add_month_begin(calendar, initial_date_args, offset, expected_date_args):
516    date_type = get_date_type(calendar)
517    initial = date_type(*initial_date_args)
518    result = initial + offset
519    expected = date_type(*expected_date_args)
520    assert result == expected
521
522
523@pytest.mark.parametrize(
524    ("initial_date_args", "offset", "expected_year_month", "expected_sub_day"),
525    [
526        ((1, 1, 1), MonthEnd(), (1, 1), ()),
527        ((1, 1, 1), MonthEnd(n=2), (1, 2), ()),
528        ((1, 3, 1), MonthEnd(n=-1), (1, 2), ()),
529        ((1, 3, 1), MonthEnd(n=-2), (1, 1), ()),
530        ((1, 2, 1), MonthEnd(n=14), (2, 3), ()),
531        ((2, 4, 1), MonthEnd(n=-14), (1, 2), ()),
532        ((1, 1, 1, 5, 5, 5, 5), MonthEnd(), (1, 1), (5, 5, 5, 5)),
533        ((1, 2, 1, 5, 5, 5, 5), MonthEnd(n=-1), (1, 1), (5, 5, 5, 5)),
534    ],
535    ids=_id_func,
536)
537def test_add_month_end(
538    calendar, initial_date_args, offset, expected_year_month, expected_sub_day
539):
540    date_type = get_date_type(calendar)
541    initial = date_type(*initial_date_args)
542    result = initial + offset
543    reference_args = expected_year_month + (1,)
544    reference = date_type(*reference_args)
545
546    # Here the days at the end of each month varies based on the calendar used
547    expected_date_args = (
548        expected_year_month + (_days_in_month(reference),) + expected_sub_day
549    )
550    expected = date_type(*expected_date_args)
551    assert result == expected
552
553
554@pytest.mark.parametrize(
555    (
556        "initial_year_month",
557        "initial_sub_day",
558        "offset",
559        "expected_year_month",
560        "expected_sub_day",
561    ),
562    [
563        ((1, 1), (), MonthEnd(), (1, 2), ()),
564        ((1, 1), (), MonthEnd(n=2), (1, 3), ()),
565        ((1, 3), (), MonthEnd(n=-1), (1, 2), ()),
566        ((1, 3), (), MonthEnd(n=-2), (1, 1), ()),
567        ((1, 2), (), MonthEnd(n=14), (2, 4), ()),
568        ((2, 4), (), MonthEnd(n=-14), (1, 2), ()),
569        ((1, 1), (5, 5, 5, 5), MonthEnd(), (1, 2), (5, 5, 5, 5)),
570        ((1, 2), (5, 5, 5, 5), MonthEnd(n=-1), (1, 1), (5, 5, 5, 5)),
571    ],
572    ids=_id_func,
573)
574def test_add_month_end_onOffset(
575    calendar,
576    initial_year_month,
577    initial_sub_day,
578    offset,
579    expected_year_month,
580    expected_sub_day,
581):
582    date_type = get_date_type(calendar)
583    reference_args = initial_year_month + (1,)
584    reference = date_type(*reference_args)
585    initial_date_args = (
586        initial_year_month + (_days_in_month(reference),) + initial_sub_day
587    )
588    initial = date_type(*initial_date_args)
589    result = initial + offset
590    reference_args = expected_year_month + (1,)
591    reference = date_type(*reference_args)
592
593    # Here the days at the end of each month varies based on the calendar used
594    expected_date_args = (
595        expected_year_month + (_days_in_month(reference),) + expected_sub_day
596    )
597    expected = date_type(*expected_date_args)
598    assert result == expected
599
600
601@pytest.mark.parametrize(
602    ("initial_date_args", "offset", "expected_date_args"),
603    [
604        ((1, 1, 1), YearBegin(), (2, 1, 1)),
605        ((1, 1, 1), YearBegin(n=2), (3, 1, 1)),
606        ((1, 1, 1), YearBegin(month=2), (1, 2, 1)),
607        ((1, 1, 7), YearBegin(n=2), (3, 1, 1)),
608        ((2, 2, 1), YearBegin(n=-1), (2, 1, 1)),
609        ((1, 1, 2), YearBegin(n=-1), (1, 1, 1)),
610        ((1, 1, 1, 5, 5, 5, 5), YearBegin(), (2, 1, 1, 5, 5, 5, 5)),
611        ((2, 1, 1, 5, 5, 5, 5), YearBegin(n=-1), (1, 1, 1, 5, 5, 5, 5)),
612    ],
613    ids=_id_func,
614)
615def test_add_year_begin(calendar, initial_date_args, offset, expected_date_args):
616    date_type = get_date_type(calendar)
617    initial = date_type(*initial_date_args)
618    result = initial + offset
619    expected = date_type(*expected_date_args)
620    assert result == expected
621
622
623@pytest.mark.parametrize(
624    ("initial_date_args", "offset", "expected_year_month", "expected_sub_day"),
625    [
626        ((1, 1, 1), YearEnd(), (1, 12), ()),
627        ((1, 1, 1), YearEnd(n=2), (2, 12), ()),
628        ((1, 1, 1), YearEnd(month=1), (1, 1), ()),
629        ((2, 3, 1), YearEnd(n=-1), (1, 12), ()),
630        ((1, 3, 1), YearEnd(n=-1, month=2), (1, 2), ()),
631        ((1, 1, 1, 5, 5, 5, 5), YearEnd(), (1, 12), (5, 5, 5, 5)),
632        ((1, 1, 1, 5, 5, 5, 5), YearEnd(n=2), (2, 12), (5, 5, 5, 5)),
633    ],
634    ids=_id_func,
635)
636def test_add_year_end(
637    calendar, initial_date_args, offset, expected_year_month, expected_sub_day
638):
639    date_type = get_date_type(calendar)
640    initial = date_type(*initial_date_args)
641    result = initial + offset
642    reference_args = expected_year_month + (1,)
643    reference = date_type(*reference_args)
644
645    # Here the days at the end of each month varies based on the calendar used
646    expected_date_args = (
647        expected_year_month + (_days_in_month(reference),) + expected_sub_day
648    )
649    expected = date_type(*expected_date_args)
650    assert result == expected
651
652
653@pytest.mark.parametrize(
654    (
655        "initial_year_month",
656        "initial_sub_day",
657        "offset",
658        "expected_year_month",
659        "expected_sub_day",
660    ),
661    [
662        ((1, 12), (), YearEnd(), (2, 12), ()),
663        ((1, 12), (), YearEnd(n=2), (3, 12), ()),
664        ((2, 12), (), YearEnd(n=-1), (1, 12), ()),
665        ((3, 12), (), YearEnd(n=-2), (1, 12), ()),
666        ((1, 1), (), YearEnd(month=2), (1, 2), ()),
667        ((1, 12), (5, 5, 5, 5), YearEnd(), (2, 12), (5, 5, 5, 5)),
668        ((2, 12), (5, 5, 5, 5), YearEnd(n=-1), (1, 12), (5, 5, 5, 5)),
669    ],
670    ids=_id_func,
671)
672def test_add_year_end_onOffset(
673    calendar,
674    initial_year_month,
675    initial_sub_day,
676    offset,
677    expected_year_month,
678    expected_sub_day,
679):
680    date_type = get_date_type(calendar)
681    reference_args = initial_year_month + (1,)
682    reference = date_type(*reference_args)
683    initial_date_args = (
684        initial_year_month + (_days_in_month(reference),) + initial_sub_day
685    )
686    initial = date_type(*initial_date_args)
687    result = initial + offset
688    reference_args = expected_year_month + (1,)
689    reference = date_type(*reference_args)
690
691    # Here the days at the end of each month varies based on the calendar used
692    expected_date_args = (
693        expected_year_month + (_days_in_month(reference),) + expected_sub_day
694    )
695    expected = date_type(*expected_date_args)
696    assert result == expected
697
698
699@pytest.mark.parametrize(
700    ("initial_date_args", "offset", "expected_date_args"),
701    [
702        ((1, 1, 1), QuarterBegin(), (1, 3, 1)),
703        ((1, 1, 1), QuarterBegin(n=2), (1, 6, 1)),
704        ((1, 1, 1), QuarterBegin(month=2), (1, 2, 1)),
705        ((1, 1, 7), QuarterBegin(n=2), (1, 6, 1)),
706        ((2, 2, 1), QuarterBegin(n=-1), (1, 12, 1)),
707        ((1, 3, 2), QuarterBegin(n=-1), (1, 3, 1)),
708        ((1, 1, 1, 5, 5, 5, 5), QuarterBegin(), (1, 3, 1, 5, 5, 5, 5)),
709        ((2, 1, 1, 5, 5, 5, 5), QuarterBegin(n=-1), (1, 12, 1, 5, 5, 5, 5)),
710    ],
711    ids=_id_func,
712)
713def test_add_quarter_begin(calendar, initial_date_args, offset, expected_date_args):
714    date_type = get_date_type(calendar)
715    initial = date_type(*initial_date_args)
716    result = initial + offset
717    expected = date_type(*expected_date_args)
718    assert result == expected
719
720
721@pytest.mark.parametrize(
722    ("initial_date_args", "offset", "expected_year_month", "expected_sub_day"),
723    [
724        ((1, 1, 1), QuarterEnd(), (1, 3), ()),
725        ((1, 1, 1), QuarterEnd(n=2), (1, 6), ()),
726        ((1, 1, 1), QuarterEnd(month=1), (1, 1), ()),
727        ((2, 3, 1), QuarterEnd(n=-1), (1, 12), ()),
728        ((1, 3, 1), QuarterEnd(n=-1, month=2), (1, 2), ()),
729        ((1, 1, 1, 5, 5, 5, 5), QuarterEnd(), (1, 3), (5, 5, 5, 5)),
730        ((1, 1, 1, 5, 5, 5, 5), QuarterEnd(n=2), (1, 6), (5, 5, 5, 5)),
731    ],
732    ids=_id_func,
733)
734def test_add_quarter_end(
735    calendar, initial_date_args, offset, expected_year_month, expected_sub_day
736):
737    date_type = get_date_type(calendar)
738    initial = date_type(*initial_date_args)
739    result = initial + offset
740    reference_args = expected_year_month + (1,)
741    reference = date_type(*reference_args)
742
743    # Here the days at the end of each month varies based on the calendar used
744    expected_date_args = (
745        expected_year_month + (_days_in_month(reference),) + expected_sub_day
746    )
747    expected = date_type(*expected_date_args)
748    assert result == expected
749
750
751@pytest.mark.parametrize(
752    (
753        "initial_year_month",
754        "initial_sub_day",
755        "offset",
756        "expected_year_month",
757        "expected_sub_day",
758    ),
759    [
760        ((1, 12), (), QuarterEnd(), (2, 3), ()),
761        ((1, 12), (), QuarterEnd(n=2), (2, 6), ()),
762        ((1, 12), (), QuarterEnd(n=-1), (1, 9), ()),
763        ((1, 12), (), QuarterEnd(n=-2), (1, 6), ()),
764        ((1, 1), (), QuarterEnd(month=2), (1, 2), ()),
765        ((1, 12), (5, 5, 5, 5), QuarterEnd(), (2, 3), (5, 5, 5, 5)),
766        ((1, 12), (5, 5, 5, 5), QuarterEnd(n=-1), (1, 9), (5, 5, 5, 5)),
767    ],
768    ids=_id_func,
769)
770def test_add_quarter_end_onOffset(
771    calendar,
772    initial_year_month,
773    initial_sub_day,
774    offset,
775    expected_year_month,
776    expected_sub_day,
777):
778    date_type = get_date_type(calendar)
779    reference_args = initial_year_month + (1,)
780    reference = date_type(*reference_args)
781    initial_date_args = (
782        initial_year_month + (_days_in_month(reference),) + initial_sub_day
783    )
784    initial = date_type(*initial_date_args)
785    result = initial + offset
786    reference_args = expected_year_month + (1,)
787    reference = date_type(*reference_args)
788
789    # Here the days at the end of each month varies based on the calendar used
790    expected_date_args = (
791        expected_year_month + (_days_in_month(reference),) + expected_sub_day
792    )
793    expected = date_type(*expected_date_args)
794    assert result == expected
795
796
797# Note for all sub-monthly offsets, pandas always returns True for onOffset
798@pytest.mark.parametrize(
799    ("date_args", "offset", "expected"),
800    [
801        ((1, 1, 1), MonthBegin(), True),
802        ((1, 1, 1, 1), MonthBegin(), True),
803        ((1, 1, 5), MonthBegin(), False),
804        ((1, 1, 5), MonthEnd(), False),
805        ((1, 3, 1), QuarterBegin(), True),
806        ((1, 3, 1, 1), QuarterBegin(), True),
807        ((1, 3, 5), QuarterBegin(), False),
808        ((1, 12, 1), QuarterEnd(), False),
809        ((1, 1, 1), YearBegin(), True),
810        ((1, 1, 1, 1), YearBegin(), True),
811        ((1, 1, 5), YearBegin(), False),
812        ((1, 12, 1), YearEnd(), False),
813        ((1, 1, 1), Day(), True),
814        ((1, 1, 1, 1), Day(), True),
815        ((1, 1, 1), Hour(), True),
816        ((1, 1, 1), Minute(), True),
817        ((1, 1, 1), Second(), True),
818        ((1, 1, 1), Millisecond(), True),
819        ((1, 1, 1), Microsecond(), True),
820    ],
821    ids=_id_func,
822)
823def test_onOffset(calendar, date_args, offset, expected):
824    date_type = get_date_type(calendar)
825    date = date_type(*date_args)
826    result = offset.onOffset(date)
827    assert result == expected
828
829
830@pytest.mark.parametrize(
831    ("year_month_args", "sub_day_args", "offset"),
832    [
833        ((1, 1), (), MonthEnd()),
834        ((1, 1), (1,), MonthEnd()),
835        ((1, 12), (), QuarterEnd()),
836        ((1, 1), (), QuarterEnd(month=1)),
837        ((1, 12), (), YearEnd()),
838        ((1, 1), (), YearEnd(month=1)),
839    ],
840    ids=_id_func,
841)
842def test_onOffset_month_or_quarter_or_year_end(
843    calendar, year_month_args, sub_day_args, offset
844):
845    date_type = get_date_type(calendar)
846    reference_args = year_month_args + (1,)
847    reference = date_type(*reference_args)
848    date_args = year_month_args + (_days_in_month(reference),) + sub_day_args
849    date = date_type(*date_args)
850    result = offset.onOffset(date)
851    assert result
852
853
854@pytest.mark.parametrize(
855    ("offset", "initial_date_args", "partial_expected_date_args"),
856    [
857        (YearBegin(), (1, 3, 1), (2, 1)),
858        (YearBegin(), (1, 1, 1), (1, 1)),
859        (YearBegin(n=2), (1, 3, 1), (2, 1)),
860        (YearBegin(n=2, month=2), (1, 3, 1), (2, 2)),
861        (YearEnd(), (1, 3, 1), (1, 12)),
862        (YearEnd(n=2), (1, 3, 1), (1, 12)),
863        (YearEnd(n=2, month=2), (1, 3, 1), (2, 2)),
864        (YearEnd(n=2, month=4), (1, 4, 30), (1, 4)),
865        (QuarterBegin(), (1, 3, 2), (1, 6)),
866        (QuarterBegin(), (1, 4, 1), (1, 6)),
867        (QuarterBegin(n=2), (1, 4, 1), (1, 6)),
868        (QuarterBegin(n=2, month=2), (1, 4, 1), (1, 5)),
869        (QuarterEnd(), (1, 3, 1), (1, 3)),
870        (QuarterEnd(n=2), (1, 3, 1), (1, 3)),
871        (QuarterEnd(n=2, month=2), (1, 3, 1), (1, 5)),
872        (QuarterEnd(n=2, month=4), (1, 4, 30), (1, 4)),
873        (MonthBegin(), (1, 3, 2), (1, 4)),
874        (MonthBegin(), (1, 3, 1), (1, 3)),
875        (MonthBegin(n=2), (1, 3, 2), (1, 4)),
876        (MonthEnd(), (1, 3, 2), (1, 3)),
877        (MonthEnd(), (1, 4, 30), (1, 4)),
878        (MonthEnd(n=2), (1, 3, 2), (1, 3)),
879        (Day(), (1, 3, 2, 1), (1, 3, 2, 1)),
880        (Hour(), (1, 3, 2, 1, 1), (1, 3, 2, 1, 1)),
881        (Minute(), (1, 3, 2, 1, 1, 1), (1, 3, 2, 1, 1, 1)),
882        (Second(), (1, 3, 2, 1, 1, 1, 1), (1, 3, 2, 1, 1, 1, 1)),
883        (Millisecond(), (1, 3, 2, 1, 1, 1, 1000), (1, 3, 2, 1, 1, 1, 1000)),
884        (Microsecond(), (1, 3, 2, 1, 1, 1, 1), (1, 3, 2, 1, 1, 1, 1)),
885    ],
886    ids=_id_func,
887)
888def test_rollforward(calendar, offset, initial_date_args, partial_expected_date_args):
889    date_type = get_date_type(calendar)
890    initial = date_type(*initial_date_args)
891    if isinstance(offset, (MonthBegin, QuarterBegin, YearBegin)):
892        expected_date_args = partial_expected_date_args + (1,)
893    elif isinstance(offset, (MonthEnd, QuarterEnd, YearEnd)):
894        reference_args = partial_expected_date_args + (1,)
895        reference = date_type(*reference_args)
896        expected_date_args = partial_expected_date_args + (_days_in_month(reference),)
897    else:
898        expected_date_args = partial_expected_date_args
899    expected = date_type(*expected_date_args)
900    result = offset.rollforward(initial)
901    assert result == expected
902
903
904@pytest.mark.parametrize(
905    ("offset", "initial_date_args", "partial_expected_date_args"),
906    [
907        (YearBegin(), (1, 3, 1), (1, 1)),
908        (YearBegin(n=2), (1, 3, 1), (1, 1)),
909        (YearBegin(n=2, month=2), (1, 3, 1), (1, 2)),
910        (YearBegin(), (1, 1, 1), (1, 1)),
911        (YearBegin(n=2, month=2), (1, 2, 1), (1, 2)),
912        (YearEnd(), (2, 3, 1), (1, 12)),
913        (YearEnd(n=2), (2, 3, 1), (1, 12)),
914        (YearEnd(n=2, month=2), (2, 3, 1), (2, 2)),
915        (YearEnd(month=4), (1, 4, 30), (1, 4)),
916        (QuarterBegin(), (1, 3, 2), (1, 3)),
917        (QuarterBegin(), (1, 4, 1), (1, 3)),
918        (QuarterBegin(n=2), (1, 4, 1), (1, 3)),
919        (QuarterBegin(n=2, month=2), (1, 4, 1), (1, 2)),
920        (QuarterEnd(), (2, 3, 1), (1, 12)),
921        (QuarterEnd(n=2), (2, 3, 1), (1, 12)),
922        (QuarterEnd(n=2, month=2), (2, 3, 1), (2, 2)),
923        (QuarterEnd(n=2, month=4), (1, 4, 30), (1, 4)),
924        (MonthBegin(), (1, 3, 2), (1, 3)),
925        (MonthBegin(n=2), (1, 3, 2), (1, 3)),
926        (MonthBegin(), (1, 3, 1), (1, 3)),
927        (MonthEnd(), (1, 3, 2), (1, 2)),
928        (MonthEnd(n=2), (1, 3, 2), (1, 2)),
929        (MonthEnd(), (1, 4, 30), (1, 4)),
930        (Day(), (1, 3, 2, 1), (1, 3, 2, 1)),
931        (Hour(), (1, 3, 2, 1, 1), (1, 3, 2, 1, 1)),
932        (Minute(), (1, 3, 2, 1, 1, 1), (1, 3, 2, 1, 1, 1)),
933        (Second(), (1, 3, 2, 1, 1, 1, 1), (1, 3, 2, 1, 1, 1, 1)),
934        (Millisecond(), (1, 3, 2, 1, 1, 1, 1000), (1, 3, 2, 1, 1, 1, 1000)),
935        (Microsecond(), (1, 3, 2, 1, 1, 1, 1), (1, 3, 2, 1, 1, 1, 1)),
936    ],
937    ids=_id_func,
938)
939def test_rollback(calendar, offset, initial_date_args, partial_expected_date_args):
940    date_type = get_date_type(calendar)
941    initial = date_type(*initial_date_args)
942    if isinstance(offset, (MonthBegin, QuarterBegin, YearBegin)):
943        expected_date_args = partial_expected_date_args + (1,)
944    elif isinstance(offset, (MonthEnd, QuarterEnd, YearEnd)):
945        reference_args = partial_expected_date_args + (1,)
946        reference = date_type(*reference_args)
947        expected_date_args = partial_expected_date_args + (_days_in_month(reference),)
948    else:
949        expected_date_args = partial_expected_date_args
950    expected = date_type(*expected_date_args)
951    result = offset.rollback(initial)
952    assert result == expected
953
954
955_CFTIME_RANGE_TESTS = [
956    (
957        "0001-01-01",
958        "0001-01-04",
959        None,
960        "D",
961        None,
962        False,
963        [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4)],
964    ),
965    (
966        "0001-01-01",
967        "0001-01-04",
968        None,
969        "D",
970        "left",
971        False,
972        [(1, 1, 1), (1, 1, 2), (1, 1, 3)],
973    ),
974    (
975        "0001-01-01",
976        "0001-01-04",
977        None,
978        "D",
979        "right",
980        False,
981        [(1, 1, 2), (1, 1, 3), (1, 1, 4)],
982    ),
983    (
984        "0001-01-01T01:00:00",
985        "0001-01-04",
986        None,
987        "D",
988        None,
989        False,
990        [(1, 1, 1, 1), (1, 1, 2, 1), (1, 1, 3, 1)],
991    ),
992    (
993        "0001-01-01 01:00:00",
994        "0001-01-04",
995        None,
996        "D",
997        None,
998        False,
999        [(1, 1, 1, 1), (1, 1, 2, 1), (1, 1, 3, 1)],
1000    ),
1001    (
1002        "0001-01-01T01:00:00",
1003        "0001-01-04",
1004        None,
1005        "D",
1006        None,
1007        True,
1008        [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4)],
1009    ),
1010    (
1011        "0001-01-01",
1012        None,
1013        4,
1014        "D",
1015        None,
1016        False,
1017        [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4)],
1018    ),
1019    (
1020        None,
1021        "0001-01-04",
1022        4,
1023        "D",
1024        None,
1025        False,
1026        [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4)],
1027    ),
1028    (
1029        (1, 1, 1),
1030        "0001-01-04",
1031        None,
1032        "D",
1033        None,
1034        False,
1035        [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4)],
1036    ),
1037    (
1038        (1, 1, 1),
1039        (1, 1, 4),
1040        None,
1041        "D",
1042        None,
1043        False,
1044        [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4)],
1045    ),
1046    (
1047        "0001-01-30",
1048        "0011-02-01",
1049        None,
1050        "3AS-JUN",
1051        None,
1052        False,
1053        [(1, 6, 1), (4, 6, 1), (7, 6, 1), (10, 6, 1)],
1054    ),
1055    ("0001-01-04", "0001-01-01", None, "D", None, False, []),
1056    (
1057        "0010",
1058        None,
1059        4,
1060        YearBegin(n=-2),
1061        None,
1062        False,
1063        [(10, 1, 1), (8, 1, 1), (6, 1, 1), (4, 1, 1)],
1064    ),
1065    (
1066        "0001-01-01",
1067        "0001-01-04",
1068        4,
1069        None,
1070        None,
1071        False,
1072        [(1, 1, 1), (1, 1, 2), (1, 1, 3), (1, 1, 4)],
1073    ),
1074    (
1075        "0001-06-01",
1076        None,
1077        4,
1078        "3QS-JUN",
1079        None,
1080        False,
1081        [(1, 6, 1), (2, 3, 1), (2, 12, 1), (3, 9, 1)],
1082    ),
1083]
1084
1085
1086@pytest.mark.parametrize(
1087    ("start", "end", "periods", "freq", "closed", "normalize", "expected_date_args"),
1088    _CFTIME_RANGE_TESTS,
1089    ids=_id_func,
1090)
1091def test_cftime_range(
1092    start, end, periods, freq, closed, normalize, calendar, expected_date_args
1093):
1094    date_type = get_date_type(calendar)
1095    expected_dates = [date_type(*args) for args in expected_date_args]
1096
1097    if isinstance(start, tuple):
1098        start = date_type(*start)
1099    if isinstance(end, tuple):
1100        end = date_type(*end)
1101
1102    result = cftime_range(
1103        start=start,
1104        end=end,
1105        periods=periods,
1106        freq=freq,
1107        closed=closed,
1108        normalize=normalize,
1109        calendar=calendar,
1110    )
1111    resulting_dates = result.values
1112
1113    assert isinstance(result, CFTimeIndex)
1114
1115    if freq is not None:
1116        np.testing.assert_equal(resulting_dates, expected_dates)
1117    else:
1118        # If we create a linear range of dates using cftime.num2date
1119        # we will not get exact round number dates.  This is because
1120        # datetime arithmetic in cftime is accurate approximately to
1121        # 1 millisecond (see https://unidata.github.io/cftime/api.html).
1122        deltas = resulting_dates - expected_dates
1123        deltas = np.array([delta.total_seconds() for delta in deltas])
1124        assert np.max(np.abs(deltas)) < 0.001
1125
1126
1127def test_cftime_range_name():
1128    result = cftime_range(start="2000", periods=4, name="foo")
1129    assert result.name == "foo"
1130
1131    result = cftime_range(start="2000", periods=4)
1132    assert result.name is None
1133
1134
1135@pytest.mark.parametrize(
1136    ("start", "end", "periods", "freq", "closed"),
1137    [
1138        (None, None, 5, "A", None),
1139        ("2000", None, None, "A", None),
1140        (None, "2000", None, "A", None),
1141        ("2000", "2001", None, None, None),
1142        (None, None, None, None, None),
1143        ("2000", "2001", None, "A", "up"),
1144        ("2000", "2001", 5, "A", None),
1145    ],
1146)
1147def test_invalid_cftime_range_inputs(start, end, periods, freq, closed):
1148    with pytest.raises(ValueError):
1149        cftime_range(start, end, periods, freq, closed=closed)
1150
1151
1152_CALENDAR_SPECIFIC_MONTH_END_TESTS = [
1153    ("2M", "noleap", [(2, 28), (4, 30), (6, 30), (8, 31), (10, 31), (12, 31)]),
1154    ("2M", "all_leap", [(2, 29), (4, 30), (6, 30), (8, 31), (10, 31), (12, 31)]),
1155    ("2M", "360_day", [(2, 30), (4, 30), (6, 30), (8, 30), (10, 30), (12, 30)]),
1156    ("2M", "standard", [(2, 29), (4, 30), (6, 30), (8, 31), (10, 31), (12, 31)]),
1157    ("2M", "gregorian", [(2, 29), (4, 30), (6, 30), (8, 31), (10, 31), (12, 31)]),
1158    ("2M", "julian", [(2, 29), (4, 30), (6, 30), (8, 31), (10, 31), (12, 31)]),
1159]
1160
1161
1162@pytest.mark.parametrize(
1163    ("freq", "calendar", "expected_month_day"),
1164    _CALENDAR_SPECIFIC_MONTH_END_TESTS,
1165    ids=_id_func,
1166)
1167def test_calendar_specific_month_end(freq, calendar, expected_month_day):
1168    year = 2000  # Use a leap-year to highlight calendar differences
1169    result = cftime_range(
1170        start="2000-02", end="2001", freq=freq, calendar=calendar
1171    ).values
1172    date_type = get_date_type(calendar)
1173    expected = [date_type(year, *args) for args in expected_month_day]
1174    np.testing.assert_equal(result, expected)
1175
1176
1177@pytest.mark.parametrize(
1178    ("calendar", "start", "end", "expected_number_of_days"),
1179    [
1180        ("noleap", "2000", "2001", 365),
1181        ("all_leap", "2000", "2001", 366),
1182        ("360_day", "2000", "2001", 360),
1183        ("standard", "2000", "2001", 366),
1184        ("gregorian", "2000", "2001", 366),
1185        ("julian", "2000", "2001", 366),
1186        ("noleap", "2001", "2002", 365),
1187        ("all_leap", "2001", "2002", 366),
1188        ("360_day", "2001", "2002", 360),
1189        ("standard", "2001", "2002", 365),
1190        ("gregorian", "2001", "2002", 365),
1191        ("julian", "2001", "2002", 365),
1192    ],
1193)
1194def test_calendar_year_length(calendar, start, end, expected_number_of_days):
1195    result = cftime_range(start, end, freq="D", closed="left", calendar=calendar)
1196    assert len(result) == expected_number_of_days
1197
1198
1199@pytest.mark.parametrize("freq", ["A", "M", "D"])
1200def test_dayofweek_after_cftime_range(freq):
1201    pytest.importorskip("cftime", minversion="1.0.2.1")
1202    result = cftime_range("2000-02-01", periods=3, freq=freq).dayofweek
1203    expected = pd.date_range("2000-02-01", periods=3, freq=freq).dayofweek
1204    np.testing.assert_array_equal(result, expected)
1205
1206
1207@pytest.mark.parametrize("freq", ["A", "M", "D"])
1208def test_dayofyear_after_cftime_range(freq):
1209    pytest.importorskip("cftime", minversion="1.0.2.1")
1210    result = cftime_range("2000-02-01", periods=3, freq=freq).dayofyear
1211    expected = pd.date_range("2000-02-01", periods=3, freq=freq).dayofyear
1212    np.testing.assert_array_equal(result, expected)
1213
1214
1215def test_cftime_range_standard_calendar_refers_to_gregorian():
1216    from cftime import DatetimeGregorian
1217
1218    (result,) = cftime_range("2000", periods=1)
1219    assert isinstance(result, DatetimeGregorian)
1220