1from datetime import timedelta 2from distutils.version import LooseVersion 3from textwrap import dedent 4 5import numpy as np 6import pandas as pd 7import pytest 8 9import xarray as xr 10from xarray.coding.cftimeindex import ( 11 CFTimeIndex, 12 _parse_array_of_cftime_strings, 13 _parse_iso8601_with_reso, 14 _parsed_string_to_bounds, 15 assert_all_valid_date_type, 16 parse_iso8601_like, 17) 18from xarray.tests import assert_array_equal, assert_identical 19 20from . import requires_cftime 21from .test_coding_times import ( 22 _ALL_CALENDARS, 23 _NON_STANDARD_CALENDARS, 24 _all_cftime_date_types, 25) 26 27 28def date_dict(year=None, month=None, day=None, hour=None, minute=None, second=None): 29 return dict( 30 year=year, month=month, day=day, hour=hour, minute=minute, second=second 31 ) 32 33 34ISO8601_LIKE_STRING_TESTS = { 35 "year": ("1999", date_dict(year="1999")), 36 "month": ("199901", date_dict(year="1999", month="01")), 37 "month-dash": ("1999-01", date_dict(year="1999", month="01")), 38 "day": ("19990101", date_dict(year="1999", month="01", day="01")), 39 "day-dash": ("1999-01-01", date_dict(year="1999", month="01", day="01")), 40 "hour": ("19990101T12", date_dict(year="1999", month="01", day="01", hour="12")), 41 "hour-dash": ( 42 "1999-01-01T12", 43 date_dict(year="1999", month="01", day="01", hour="12"), 44 ), 45 "hour-space-separator": ( 46 "1999-01-01 12", 47 date_dict(year="1999", month="01", day="01", hour="12"), 48 ), 49 "minute": ( 50 "19990101T1234", 51 date_dict(year="1999", month="01", day="01", hour="12", minute="34"), 52 ), 53 "minute-dash": ( 54 "1999-01-01T12:34", 55 date_dict(year="1999", month="01", day="01", hour="12", minute="34"), 56 ), 57 "minute-space-separator": ( 58 "1999-01-01 12:34", 59 date_dict(year="1999", month="01", day="01", hour="12", minute="34"), 60 ), 61 "second": ( 62 "19990101T123456", 63 date_dict( 64 year="1999", month="01", day="01", hour="12", minute="34", second="56" 65 ), 66 ), 67 "second-dash": ( 68 "1999-01-01T12:34:56", 69 date_dict( 70 year="1999", month="01", day="01", hour="12", minute="34", second="56" 71 ), 72 ), 73 "second-space-separator": ( 74 "1999-01-01 12:34:56", 75 date_dict( 76 year="1999", month="01", day="01", hour="12", minute="34", second="56" 77 ), 78 ), 79} 80 81 82@pytest.mark.parametrize( 83 ("string", "expected"), 84 list(ISO8601_LIKE_STRING_TESTS.values()), 85 ids=list(ISO8601_LIKE_STRING_TESTS.keys()), 86) 87def test_parse_iso8601_like(string, expected): 88 result = parse_iso8601_like(string) 89 assert result == expected 90 91 with pytest.raises(ValueError): 92 parse_iso8601_like(string + "3") 93 parse_iso8601_like(string + ".3") 94 95 96_CFTIME_CALENDARS = [ 97 "365_day", 98 "360_day", 99 "julian", 100 "all_leap", 101 "366_day", 102 "gregorian", 103 "proleptic_gregorian", 104] 105 106 107@pytest.fixture(params=_CFTIME_CALENDARS) 108def date_type(request): 109 return _all_cftime_date_types()[request.param] 110 111 112@pytest.fixture 113def index(date_type): 114 dates = [ 115 date_type(1, 1, 1), 116 date_type(1, 2, 1), 117 date_type(2, 1, 1), 118 date_type(2, 2, 1), 119 ] 120 return CFTimeIndex(dates) 121 122 123@pytest.fixture 124def monotonic_decreasing_index(date_type): 125 dates = [ 126 date_type(2, 2, 1), 127 date_type(2, 1, 1), 128 date_type(1, 2, 1), 129 date_type(1, 1, 1), 130 ] 131 return CFTimeIndex(dates) 132 133 134@pytest.fixture 135def length_one_index(date_type): 136 dates = [date_type(1, 1, 1)] 137 return CFTimeIndex(dates) 138 139 140@pytest.fixture 141def da(index): 142 return xr.DataArray([1, 2, 3, 4], coords=[index], dims=["time"]) 143 144 145@pytest.fixture 146def series(index): 147 return pd.Series([1, 2, 3, 4], index=index) 148 149 150@pytest.fixture 151def df(index): 152 return pd.DataFrame([1, 2, 3, 4], index=index) 153 154 155@pytest.fixture 156def feb_days(date_type): 157 import cftime 158 159 if date_type is cftime.DatetimeAllLeap: 160 return 29 161 elif date_type is cftime.Datetime360Day: 162 return 30 163 else: 164 return 28 165 166 167@pytest.fixture 168def dec_days(date_type): 169 import cftime 170 171 if date_type is cftime.Datetime360Day: 172 return 30 173 else: 174 return 31 175 176 177@pytest.fixture 178def index_with_name(date_type): 179 dates = [ 180 date_type(1, 1, 1), 181 date_type(1, 2, 1), 182 date_type(2, 1, 1), 183 date_type(2, 2, 1), 184 ] 185 return CFTimeIndex(dates, name="foo") 186 187 188@requires_cftime 189@pytest.mark.parametrize(("name", "expected_name"), [("bar", "bar"), (None, "foo")]) 190def test_constructor_with_name(index_with_name, name, expected_name): 191 result = CFTimeIndex(index_with_name, name=name).name 192 assert result == expected_name 193 194 195@requires_cftime 196def test_assert_all_valid_date_type(date_type, index): 197 import cftime 198 199 if date_type is cftime.DatetimeNoLeap: 200 mixed_date_types = np.array( 201 [date_type(1, 1, 1), cftime.DatetimeAllLeap(1, 2, 1)] 202 ) 203 else: 204 mixed_date_types = np.array( 205 [date_type(1, 1, 1), cftime.DatetimeNoLeap(1, 2, 1)] 206 ) 207 with pytest.raises(TypeError): 208 assert_all_valid_date_type(mixed_date_types) 209 210 with pytest.raises(TypeError): 211 assert_all_valid_date_type(np.array([1, date_type(1, 1, 1)])) 212 213 assert_all_valid_date_type(np.array([date_type(1, 1, 1), date_type(1, 2, 1)])) 214 215 216@requires_cftime 217@pytest.mark.parametrize( 218 ("field", "expected"), 219 [ 220 ("year", [1, 1, 2, 2]), 221 ("month", [1, 2, 1, 2]), 222 ("day", [1, 1, 1, 1]), 223 ("hour", [0, 0, 0, 0]), 224 ("minute", [0, 0, 0, 0]), 225 ("second", [0, 0, 0, 0]), 226 ("microsecond", [0, 0, 0, 0]), 227 ], 228) 229def test_cftimeindex_field_accessors(index, field, expected): 230 result = getattr(index, field) 231 assert_array_equal(result, expected) 232 233 234@requires_cftime 235def test_cftimeindex_dayofyear_accessor(index): 236 result = index.dayofyear 237 expected = [date.dayofyr for date in index] 238 assert_array_equal(result, expected) 239 240 241@requires_cftime 242def test_cftimeindex_dayofweek_accessor(index): 243 result = index.dayofweek 244 expected = [date.dayofwk for date in index] 245 assert_array_equal(result, expected) 246 247 248@requires_cftime 249def test_cftimeindex_days_in_month_accessor(index): 250 result = index.days_in_month 251 expected = [date.daysinmonth for date in index] 252 assert_array_equal(result, expected) 253 254 255@requires_cftime 256@pytest.mark.parametrize( 257 ("string", "date_args", "reso"), 258 [ 259 ("1999", (1999, 1, 1), "year"), 260 ("199902", (1999, 2, 1), "month"), 261 ("19990202", (1999, 2, 2), "day"), 262 ("19990202T01", (1999, 2, 2, 1), "hour"), 263 ("19990202T0101", (1999, 2, 2, 1, 1), "minute"), 264 ("19990202T010156", (1999, 2, 2, 1, 1, 56), "second"), 265 ], 266) 267def test_parse_iso8601_with_reso(date_type, string, date_args, reso): 268 expected_date = date_type(*date_args) 269 expected_reso = reso 270 result_date, result_reso = _parse_iso8601_with_reso(date_type, string) 271 assert result_date == expected_date 272 assert result_reso == expected_reso 273 274 275@requires_cftime 276def test_parse_string_to_bounds_year(date_type, dec_days): 277 parsed = date_type(2, 2, 10, 6, 2, 8, 1) 278 expected_start = date_type(2, 1, 1) 279 expected_end = date_type(2, 12, dec_days, 23, 59, 59, 999999) 280 result_start, result_end = _parsed_string_to_bounds(date_type, "year", parsed) 281 assert result_start == expected_start 282 assert result_end == expected_end 283 284 285@requires_cftime 286def test_parse_string_to_bounds_month_feb(date_type, feb_days): 287 parsed = date_type(2, 2, 10, 6, 2, 8, 1) 288 expected_start = date_type(2, 2, 1) 289 expected_end = date_type(2, 2, feb_days, 23, 59, 59, 999999) 290 result_start, result_end = _parsed_string_to_bounds(date_type, "month", parsed) 291 assert result_start == expected_start 292 assert result_end == expected_end 293 294 295@requires_cftime 296def test_parse_string_to_bounds_month_dec(date_type, dec_days): 297 parsed = date_type(2, 12, 1) 298 expected_start = date_type(2, 12, 1) 299 expected_end = date_type(2, 12, dec_days, 23, 59, 59, 999999) 300 result_start, result_end = _parsed_string_to_bounds(date_type, "month", parsed) 301 assert result_start == expected_start 302 assert result_end == expected_end 303 304 305@requires_cftime 306@pytest.mark.parametrize( 307 ("reso", "ex_start_args", "ex_end_args"), 308 [ 309 ("day", (2, 2, 10), (2, 2, 10, 23, 59, 59, 999999)), 310 ("hour", (2, 2, 10, 6), (2, 2, 10, 6, 59, 59, 999999)), 311 ("minute", (2, 2, 10, 6, 2), (2, 2, 10, 6, 2, 59, 999999)), 312 ("second", (2, 2, 10, 6, 2, 8), (2, 2, 10, 6, 2, 8, 999999)), 313 ], 314) 315def test_parsed_string_to_bounds_sub_monthly( 316 date_type, reso, ex_start_args, ex_end_args 317): 318 parsed = date_type(2, 2, 10, 6, 2, 8, 123456) 319 expected_start = date_type(*ex_start_args) 320 expected_end = date_type(*ex_end_args) 321 322 result_start, result_end = _parsed_string_to_bounds(date_type, reso, parsed) 323 assert result_start == expected_start 324 assert result_end == expected_end 325 326 327@requires_cftime 328def test_parsed_string_to_bounds_raises(date_type): 329 with pytest.raises(KeyError): 330 _parsed_string_to_bounds(date_type, "a", date_type(1, 1, 1)) 331 332 333@requires_cftime 334def test_get_loc(date_type, index): 335 result = index.get_loc("0001") 336 assert result == slice(0, 2) 337 338 result = index.get_loc(date_type(1, 2, 1)) 339 assert result == 1 340 341 result = index.get_loc("0001-02-01") 342 assert result == slice(1, 2) 343 344 with pytest.raises(KeyError, match=r"1234"): 345 index.get_loc("1234") 346 347 348@requires_cftime 349def test_get_slice_bound(date_type, index): 350 # The kind argument is required in earlier versions of pandas even though it 351 # is not used by CFTimeIndex. This logic can be removed once our minimum 352 # version of pandas is at least 1.3. 353 if LooseVersion(pd.__version__) < LooseVersion("1.3"): 354 kind_args = ("getitem",) 355 else: 356 kind_args = () 357 358 result = index.get_slice_bound("0001", "left", *kind_args) 359 expected = 0 360 assert result == expected 361 362 result = index.get_slice_bound("0001", "right", *kind_args) 363 expected = 2 364 assert result == expected 365 366 result = index.get_slice_bound(date_type(1, 3, 1), "left", *kind_args) 367 expected = 2 368 assert result == expected 369 370 result = index.get_slice_bound(date_type(1, 3, 1), "right", *kind_args) 371 expected = 2 372 assert result == expected 373 374 375@requires_cftime 376def test_get_slice_bound_decreasing_index(date_type, monotonic_decreasing_index): 377 # The kind argument is required in earlier versions of pandas even though it 378 # is not used by CFTimeIndex. This logic can be removed once our minimum 379 # version of pandas is at least 1.3. 380 if LooseVersion(pd.__version__) < LooseVersion("1.3"): 381 kind_args = ("getitem",) 382 else: 383 kind_args = () 384 385 result = monotonic_decreasing_index.get_slice_bound("0001", "left", *kind_args) 386 expected = 2 387 assert result == expected 388 389 result = monotonic_decreasing_index.get_slice_bound("0001", "right", *kind_args) 390 expected = 4 391 assert result == expected 392 393 result = monotonic_decreasing_index.get_slice_bound( 394 date_type(1, 3, 1), "left", *kind_args 395 ) 396 expected = 2 397 assert result == expected 398 399 result = monotonic_decreasing_index.get_slice_bound( 400 date_type(1, 3, 1), "right", *kind_args 401 ) 402 expected = 2 403 assert result == expected 404 405 406@requires_cftime 407def test_get_slice_bound_length_one_index(date_type, length_one_index): 408 # The kind argument is required in earlier versions of pandas even though it 409 # is not used by CFTimeIndex. This logic can be removed once our minimum 410 # version of pandas is at least 1.3. 411 if LooseVersion(pd.__version__) <= LooseVersion("1.3"): 412 kind_args = ("getitem",) 413 else: 414 kind_args = () 415 416 result = length_one_index.get_slice_bound("0001", "left", *kind_args) 417 expected = 0 418 assert result == expected 419 420 result = length_one_index.get_slice_bound("0001", "right", *kind_args) 421 expected = 1 422 assert result == expected 423 424 result = length_one_index.get_slice_bound(date_type(1, 3, 1), "left", *kind_args) 425 expected = 1 426 assert result == expected 427 428 result = length_one_index.get_slice_bound(date_type(1, 3, 1), "right", *kind_args) 429 expected = 1 430 assert result == expected 431 432 433@requires_cftime 434def test_string_slice_length_one_index(length_one_index): 435 da = xr.DataArray([1], coords=[length_one_index], dims=["time"]) 436 result = da.sel(time=slice("0001", "0001")) 437 assert_identical(result, da) 438 439 440@requires_cftime 441def test_date_type_property(date_type, index): 442 assert index.date_type is date_type 443 444 445@requires_cftime 446def test_contains(date_type, index): 447 assert "0001-01-01" in index 448 assert "0001" in index 449 assert "0003" not in index 450 assert date_type(1, 1, 1) in index 451 assert date_type(3, 1, 1) not in index 452 453 454@requires_cftime 455def test_groupby(da): 456 result = da.groupby("time.month").sum("time") 457 expected = xr.DataArray([4, 6], coords=[[1, 2]], dims=["month"]) 458 assert_identical(result, expected) 459 460 461SEL_STRING_OR_LIST_TESTS = { 462 "string": "0001", 463 "string-slice": slice("0001-01-01", "0001-12-30"), 464 "bool-list": [True, True, False, False], 465} 466 467 468@requires_cftime 469@pytest.mark.parametrize( 470 "sel_arg", 471 list(SEL_STRING_OR_LIST_TESTS.values()), 472 ids=list(SEL_STRING_OR_LIST_TESTS.keys()), 473) 474def test_sel_string_or_list(da, index, sel_arg): 475 expected = xr.DataArray([1, 2], coords=[index[:2]], dims=["time"]) 476 result = da.sel(time=sel_arg) 477 assert_identical(result, expected) 478 479 480@requires_cftime 481def test_sel_date_slice_or_list(da, index, date_type): 482 expected = xr.DataArray([1, 2], coords=[index[:2]], dims=["time"]) 483 result = da.sel(time=slice(date_type(1, 1, 1), date_type(1, 12, 30))) 484 assert_identical(result, expected) 485 486 result = da.sel(time=[date_type(1, 1, 1), date_type(1, 2, 1)]) 487 assert_identical(result, expected) 488 489 490@requires_cftime 491def test_sel_date_scalar(da, date_type, index): 492 expected = xr.DataArray(1).assign_coords(time=index[0]) 493 result = da.sel(time=date_type(1, 1, 1)) 494 assert_identical(result, expected) 495 496 497@requires_cftime 498def test_sel_date_distant_date(da, date_type, index): 499 expected = xr.DataArray(4).assign_coords(time=index[3]) 500 result = da.sel(time=date_type(2000, 1, 1), method="nearest") 501 assert_identical(result, expected) 502 503 504@requires_cftime 505@pytest.mark.parametrize( 506 "sel_kwargs", 507 [ 508 {"method": "nearest"}, 509 {"method": "nearest", "tolerance": timedelta(days=70)}, 510 {"method": "nearest", "tolerance": timedelta(days=1800000)}, 511 ], 512) 513def test_sel_date_scalar_nearest(da, date_type, index, sel_kwargs): 514 expected = xr.DataArray(2).assign_coords(time=index[1]) 515 result = da.sel(time=date_type(1, 4, 1), **sel_kwargs) 516 assert_identical(result, expected) 517 518 expected = xr.DataArray(3).assign_coords(time=index[2]) 519 result = da.sel(time=date_type(1, 11, 1), **sel_kwargs) 520 assert_identical(result, expected) 521 522 523@requires_cftime 524@pytest.mark.parametrize( 525 "sel_kwargs", 526 [{"method": "pad"}, {"method": "pad", "tolerance": timedelta(days=365)}], 527) 528def test_sel_date_scalar_pad(da, date_type, index, sel_kwargs): 529 expected = xr.DataArray(2).assign_coords(time=index[1]) 530 result = da.sel(time=date_type(1, 4, 1), **sel_kwargs) 531 assert_identical(result, expected) 532 533 expected = xr.DataArray(2).assign_coords(time=index[1]) 534 result = da.sel(time=date_type(1, 11, 1), **sel_kwargs) 535 assert_identical(result, expected) 536 537 538@requires_cftime 539@pytest.mark.parametrize( 540 "sel_kwargs", 541 [{"method": "backfill"}, {"method": "backfill", "tolerance": timedelta(days=365)}], 542) 543def test_sel_date_scalar_backfill(da, date_type, index, sel_kwargs): 544 expected = xr.DataArray(3).assign_coords(time=index[2]) 545 result = da.sel(time=date_type(1, 4, 1), **sel_kwargs) 546 assert_identical(result, expected) 547 548 expected = xr.DataArray(3).assign_coords(time=index[2]) 549 result = da.sel(time=date_type(1, 11, 1), **sel_kwargs) 550 assert_identical(result, expected) 551 552 553@requires_cftime 554@pytest.mark.parametrize( 555 "sel_kwargs", 556 [ 557 {"method": "pad", "tolerance": timedelta(days=20)}, 558 {"method": "backfill", "tolerance": timedelta(days=20)}, 559 {"method": "nearest", "tolerance": timedelta(days=20)}, 560 ], 561) 562def test_sel_date_scalar_tolerance_raises(da, date_type, sel_kwargs): 563 with pytest.raises(KeyError): 564 da.sel(time=date_type(1, 5, 1), **sel_kwargs) 565 566 567@requires_cftime 568@pytest.mark.parametrize( 569 "sel_kwargs", 570 [{"method": "nearest"}, {"method": "nearest", "tolerance": timedelta(days=70)}], 571) 572def test_sel_date_list_nearest(da, date_type, index, sel_kwargs): 573 expected = xr.DataArray([2, 2], coords=[[index[1], index[1]]], dims=["time"]) 574 result = da.sel(time=[date_type(1, 3, 1), date_type(1, 4, 1)], **sel_kwargs) 575 assert_identical(result, expected) 576 577 expected = xr.DataArray([2, 3], coords=[[index[1], index[2]]], dims=["time"]) 578 result = da.sel(time=[date_type(1, 3, 1), date_type(1, 12, 1)], **sel_kwargs) 579 assert_identical(result, expected) 580 581 expected = xr.DataArray([3, 3], coords=[[index[2], index[2]]], dims=["time"]) 582 result = da.sel(time=[date_type(1, 11, 1), date_type(1, 12, 1)], **sel_kwargs) 583 assert_identical(result, expected) 584 585 586@requires_cftime 587@pytest.mark.parametrize( 588 "sel_kwargs", 589 [{"method": "pad"}, {"method": "pad", "tolerance": timedelta(days=365)}], 590) 591def test_sel_date_list_pad(da, date_type, index, sel_kwargs): 592 expected = xr.DataArray([2, 2], coords=[[index[1], index[1]]], dims=["time"]) 593 result = da.sel(time=[date_type(1, 3, 1), date_type(1, 4, 1)], **sel_kwargs) 594 assert_identical(result, expected) 595 596 597@requires_cftime 598@pytest.mark.parametrize( 599 "sel_kwargs", 600 [{"method": "backfill"}, {"method": "backfill", "tolerance": timedelta(days=365)}], 601) 602def test_sel_date_list_backfill(da, date_type, index, sel_kwargs): 603 expected = xr.DataArray([3, 3], coords=[[index[2], index[2]]], dims=["time"]) 604 result = da.sel(time=[date_type(1, 3, 1), date_type(1, 4, 1)], **sel_kwargs) 605 assert_identical(result, expected) 606 607 608@requires_cftime 609@pytest.mark.parametrize( 610 "sel_kwargs", 611 [ 612 {"method": "pad", "tolerance": timedelta(days=20)}, 613 {"method": "backfill", "tolerance": timedelta(days=20)}, 614 {"method": "nearest", "tolerance": timedelta(days=20)}, 615 ], 616) 617def test_sel_date_list_tolerance_raises(da, date_type, sel_kwargs): 618 with pytest.raises(KeyError): 619 da.sel(time=[date_type(1, 2, 1), date_type(1, 5, 1)], **sel_kwargs) 620 621 622@requires_cftime 623def test_isel(da, index): 624 expected = xr.DataArray(1).assign_coords(time=index[0]) 625 result = da.isel(time=0) 626 assert_identical(result, expected) 627 628 expected = xr.DataArray([1, 2], coords=[index[:2]], dims=["time"]) 629 result = da.isel(time=[0, 1]) 630 assert_identical(result, expected) 631 632 633@pytest.fixture 634def scalar_args(date_type): 635 return [date_type(1, 1, 1)] 636 637 638@pytest.fixture 639def range_args(date_type): 640 return [ 641 "0001", 642 slice("0001-01-01", "0001-12-30"), 643 slice(None, "0001-12-30"), 644 slice(date_type(1, 1, 1), date_type(1, 12, 30)), 645 slice(None, date_type(1, 12, 30)), 646 ] 647 648 649@requires_cftime 650def test_indexing_in_series_getitem(series, index, scalar_args, range_args): 651 for arg in scalar_args: 652 assert series[arg] == 1 653 654 expected = pd.Series([1, 2], index=index[:2]) 655 for arg in range_args: 656 assert series[arg].equals(expected) 657 658 659@requires_cftime 660def test_indexing_in_series_loc(series, index, scalar_args, range_args): 661 for arg in scalar_args: 662 assert series.loc[arg] == 1 663 664 expected = pd.Series([1, 2], index=index[:2]) 665 for arg in range_args: 666 assert series.loc[arg].equals(expected) 667 668 669@requires_cftime 670def test_indexing_in_series_iloc(series, index): 671 expected = 1 672 assert series.iloc[0] == expected 673 674 expected = pd.Series([1, 2], index=index[:2]) 675 assert series.iloc[:2].equals(expected) 676 677 678@requires_cftime 679def test_series_dropna(index): 680 series = pd.Series([0.0, 1.0, np.nan, np.nan], index=index) 681 expected = series.iloc[:2] 682 result = series.dropna() 683 assert result.equals(expected) 684 685 686@requires_cftime 687def test_indexing_in_dataframe_loc(df, index, scalar_args, range_args): 688 expected = pd.Series([1], name=index[0]) 689 for arg in scalar_args: 690 result = df.loc[arg] 691 assert result.equals(expected) 692 693 expected = pd.DataFrame([1, 2], index=index[:2]) 694 for arg in range_args: 695 result = df.loc[arg] 696 assert result.equals(expected) 697 698 699@requires_cftime 700def test_indexing_in_dataframe_iloc(df, index): 701 expected = pd.Series([1], name=index[0]) 702 result = df.iloc[0] 703 assert result.equals(expected) 704 assert result.equals(expected) 705 706 expected = pd.DataFrame([1, 2], index=index[:2]) 707 result = df.iloc[:2] 708 assert result.equals(expected) 709 710 711@requires_cftime 712def test_concat_cftimeindex(date_type): 713 da1 = xr.DataArray( 714 [1.0, 2.0], coords=[[date_type(1, 1, 1), date_type(1, 2, 1)]], dims=["time"] 715 ) 716 da2 = xr.DataArray( 717 [3.0, 4.0], coords=[[date_type(1, 3, 1), date_type(1, 4, 1)]], dims=["time"] 718 ) 719 da = xr.concat([da1, da2], dim="time") 720 721 assert isinstance(da.xindexes["time"].to_pandas_index(), CFTimeIndex) 722 723 724@requires_cftime 725def test_empty_cftimeindex(): 726 index = CFTimeIndex([]) 727 assert index.date_type is None 728 729 730@requires_cftime 731def test_cftimeindex_add(index): 732 date_type = index.date_type 733 expected_dates = [ 734 date_type(1, 1, 2), 735 date_type(1, 2, 2), 736 date_type(2, 1, 2), 737 date_type(2, 2, 2), 738 ] 739 expected = CFTimeIndex(expected_dates) 740 result = index + timedelta(days=1) 741 assert result.equals(expected) 742 assert isinstance(result, CFTimeIndex) 743 744 745@requires_cftime 746@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 747def test_cftimeindex_add_timedeltaindex(calendar): 748 a = xr.cftime_range("2000", periods=5, calendar=calendar) 749 deltas = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)]) 750 result = a + deltas 751 expected = a.shift(2, "D") 752 assert result.equals(expected) 753 assert isinstance(result, CFTimeIndex) 754 755 756@requires_cftime 757def test_cftimeindex_radd(index): 758 date_type = index.date_type 759 expected_dates = [ 760 date_type(1, 1, 2), 761 date_type(1, 2, 2), 762 date_type(2, 1, 2), 763 date_type(2, 2, 2), 764 ] 765 expected = CFTimeIndex(expected_dates) 766 result = timedelta(days=1) + index 767 assert result.equals(expected) 768 assert isinstance(result, CFTimeIndex) 769 770 771@requires_cftime 772@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 773def test_timedeltaindex_add_cftimeindex(calendar): 774 a = xr.cftime_range("2000", periods=5, calendar=calendar) 775 deltas = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)]) 776 result = deltas + a 777 expected = a.shift(2, "D") 778 assert result.equals(expected) 779 assert isinstance(result, CFTimeIndex) 780 781 782@requires_cftime 783def test_cftimeindex_sub_timedelta(index): 784 date_type = index.date_type 785 expected_dates = [ 786 date_type(1, 1, 2), 787 date_type(1, 2, 2), 788 date_type(2, 1, 2), 789 date_type(2, 2, 2), 790 ] 791 expected = CFTimeIndex(expected_dates) 792 result = index + timedelta(days=2) 793 result = result - timedelta(days=1) 794 assert result.equals(expected) 795 assert isinstance(result, CFTimeIndex) 796 797 798@requires_cftime 799@pytest.mark.parametrize( 800 "other", 801 [np.array(4 * [timedelta(days=1)]), np.array(timedelta(days=1))], 802 ids=["1d-array", "scalar-array"], 803) 804def test_cftimeindex_sub_timedelta_array(index, other): 805 date_type = index.date_type 806 expected_dates = [ 807 date_type(1, 1, 2), 808 date_type(1, 2, 2), 809 date_type(2, 1, 2), 810 date_type(2, 2, 2), 811 ] 812 expected = CFTimeIndex(expected_dates) 813 result = index + timedelta(days=2) 814 result = result - other 815 assert result.equals(expected) 816 assert isinstance(result, CFTimeIndex) 817 818 819@requires_cftime 820@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 821def test_cftimeindex_sub_cftimeindex(calendar): 822 a = xr.cftime_range("2000", periods=5, calendar=calendar) 823 b = a.shift(2, "D") 824 result = b - a 825 expected = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)]) 826 assert result.equals(expected) 827 assert isinstance(result, pd.TimedeltaIndex) 828 829 830@requires_cftime 831@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 832def test_cftimeindex_sub_cftime_datetime(calendar): 833 a = xr.cftime_range("2000", periods=5, calendar=calendar) 834 result = a - a[0] 835 expected = pd.TimedeltaIndex([timedelta(days=i) for i in range(5)]) 836 assert result.equals(expected) 837 assert isinstance(result, pd.TimedeltaIndex) 838 839 840@requires_cftime 841@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 842def test_cftime_datetime_sub_cftimeindex(calendar): 843 a = xr.cftime_range("2000", periods=5, calendar=calendar) 844 result = a[0] - a 845 expected = pd.TimedeltaIndex([timedelta(days=-i) for i in range(5)]) 846 assert result.equals(expected) 847 assert isinstance(result, pd.TimedeltaIndex) 848 849 850@requires_cftime 851@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 852def test_distant_cftime_datetime_sub_cftimeindex(calendar): 853 a = xr.cftime_range("2000", periods=5, calendar=calendar) 854 with pytest.raises(ValueError, match="difference exceeds"): 855 a.date_type(1, 1, 1) - a 856 857 858@requires_cftime 859@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 860def test_cftimeindex_sub_timedeltaindex(calendar): 861 a = xr.cftime_range("2000", periods=5, calendar=calendar) 862 deltas = pd.TimedeltaIndex([timedelta(days=2) for _ in range(5)]) 863 result = a - deltas 864 expected = a.shift(-2, "D") 865 assert result.equals(expected) 866 assert isinstance(result, CFTimeIndex) 867 868 869@requires_cftime 870@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 871def test_cftimeindex_sub_index_of_cftime_datetimes(calendar): 872 a = xr.cftime_range("2000", periods=5, calendar=calendar) 873 b = pd.Index(a.values) 874 expected = a - a 875 result = a - b 876 assert result.equals(expected) 877 assert isinstance(result, pd.TimedeltaIndex) 878 879 880@requires_cftime 881@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 882def test_cftimeindex_sub_not_implemented(calendar): 883 a = xr.cftime_range("2000", periods=5, calendar=calendar) 884 with pytest.raises(TypeError, match="unsupported operand"): 885 a - 1 886 887 888@requires_cftime 889def test_cftimeindex_rsub(index): 890 with pytest.raises(TypeError): 891 timedelta(days=1) - index 892 893 894@requires_cftime 895@pytest.mark.parametrize("freq", ["D", timedelta(days=1)]) 896def test_cftimeindex_shift(index, freq): 897 date_type = index.date_type 898 expected_dates = [ 899 date_type(1, 1, 3), 900 date_type(1, 2, 3), 901 date_type(2, 1, 3), 902 date_type(2, 2, 3), 903 ] 904 expected = CFTimeIndex(expected_dates) 905 result = index.shift(2, freq) 906 assert result.equals(expected) 907 assert isinstance(result, CFTimeIndex) 908 909 910@requires_cftime 911def test_cftimeindex_shift_invalid_n(): 912 index = xr.cftime_range("2000", periods=3) 913 with pytest.raises(TypeError): 914 index.shift("a", "D") 915 916 917@requires_cftime 918def test_cftimeindex_shift_invalid_freq(): 919 index = xr.cftime_range("2000", periods=3) 920 with pytest.raises(TypeError): 921 index.shift(1, 1) 922 923 924@requires_cftime 925@pytest.mark.parametrize( 926 ("calendar", "expected"), 927 [ 928 ("noleap", "noleap"), 929 ("365_day", "noleap"), 930 ("360_day", "360_day"), 931 ("julian", "julian"), 932 ("gregorian", "gregorian"), 933 ("proleptic_gregorian", "proleptic_gregorian"), 934 ], 935) 936def test_cftimeindex_calendar_property(calendar, expected): 937 index = xr.cftime_range(start="2000", periods=3, calendar=calendar) 938 assert index.calendar == expected 939 940 941@requires_cftime 942@pytest.mark.parametrize( 943 ("calendar", "expected"), 944 [ 945 ("noleap", "noleap"), 946 ("365_day", "noleap"), 947 ("360_day", "360_day"), 948 ("julian", "julian"), 949 ("gregorian", "gregorian"), 950 ("proleptic_gregorian", "proleptic_gregorian"), 951 ], 952) 953def test_cftimeindex_calendar_repr(calendar, expected): 954 """Test that cftimeindex has calendar property in repr.""" 955 index = xr.cftime_range(start="2000", periods=3, calendar=calendar) 956 repr_str = index.__repr__() 957 assert f" calendar='{expected}'" in repr_str 958 assert "2000-01-01 00:00:00, 2000-01-02 00:00:00" in repr_str 959 960 961@requires_cftime 962@pytest.mark.parametrize("periods", [2, 40]) 963def test_cftimeindex_periods_repr(periods): 964 """Test that cftimeindex has periods property in repr.""" 965 index = xr.cftime_range(start="2000", periods=periods) 966 repr_str = index.__repr__() 967 assert f" length={periods}" in repr_str 968 969 970@requires_cftime 971@pytest.mark.parametrize("calendar", ["noleap", "360_day", "standard"]) 972@pytest.mark.parametrize("freq", ["D", "H"]) 973def test_cftimeindex_freq_in_repr(freq, calendar): 974 """Test that cftimeindex has frequency property in repr.""" 975 index = xr.cftime_range(start="2000", periods=3, freq=freq, calendar=calendar) 976 repr_str = index.__repr__() 977 assert f", freq='{freq}'" in repr_str 978 979 980@requires_cftime 981@pytest.mark.parametrize( 982 "periods,expected", 983 [ 984 ( 985 2, 986 """\ 987CFTimeIndex([2000-01-01 00:00:00, 2000-01-02 00:00:00], 988 dtype='object', length=2, calendar='gregorian', freq=None)""", 989 ), 990 ( 991 4, 992 """\ 993CFTimeIndex([2000-01-01 00:00:00, 2000-01-02 00:00:00, 2000-01-03 00:00:00, 994 2000-01-04 00:00:00], 995 dtype='object', length=4, calendar='gregorian', freq='D')""", 996 ), 997 ( 998 101, 999 """\ 1000CFTimeIndex([2000-01-01 00:00:00, 2000-01-02 00:00:00, 2000-01-03 00:00:00, 1001 2000-01-04 00:00:00, 2000-01-05 00:00:00, 2000-01-06 00:00:00, 1002 2000-01-07 00:00:00, 2000-01-08 00:00:00, 2000-01-09 00:00:00, 1003 2000-01-10 00:00:00, 1004 ... 1005 2000-04-01 00:00:00, 2000-04-02 00:00:00, 2000-04-03 00:00:00, 1006 2000-04-04 00:00:00, 2000-04-05 00:00:00, 2000-04-06 00:00:00, 1007 2000-04-07 00:00:00, 2000-04-08 00:00:00, 2000-04-09 00:00:00, 1008 2000-04-10 00:00:00], 1009 dtype='object', length=101, calendar='gregorian', freq='D')""", 1010 ), 1011 ], 1012) 1013def test_cftimeindex_repr_formatting(periods, expected): 1014 """Test that cftimeindex.__repr__ is formatted similar to pd.Index.__repr__.""" 1015 index = xr.cftime_range(start="2000", periods=periods, freq="D") 1016 expected = dedent(expected) 1017 assert expected == repr(index) 1018 1019 1020@requires_cftime 1021@pytest.mark.parametrize("display_width", [40, 80, 100]) 1022@pytest.mark.parametrize("periods", [2, 3, 4, 100, 101]) 1023def test_cftimeindex_repr_formatting_width(periods, display_width): 1024 """Test that cftimeindex is sensitive to OPTIONS['display_width'].""" 1025 index = xr.cftime_range(start="2000", periods=periods) 1026 len_intro_str = len("CFTimeIndex(") 1027 with xr.set_options(display_width=display_width): 1028 repr_str = index.__repr__() 1029 splitted = repr_str.split("\n") 1030 for i, s in enumerate(splitted): 1031 # check that lines not longer than OPTIONS['display_width'] 1032 assert len(s) <= display_width, f"{len(s)} {s} {display_width}" 1033 if i > 0: 1034 # check for initial spaces 1035 assert s[:len_intro_str] == " " * len_intro_str 1036 1037 1038@requires_cftime 1039@pytest.mark.parametrize("periods", [22, 50, 100]) 1040def test_cftimeindex_repr_101_shorter(periods): 1041 index_101 = xr.cftime_range(start="2000", periods=101) 1042 index_periods = xr.cftime_range(start="2000", periods=periods) 1043 index_101_repr_str = index_101.__repr__() 1044 index_periods_repr_str = index_periods.__repr__() 1045 assert len(index_101_repr_str) < len(index_periods_repr_str) 1046 1047 1048@requires_cftime 1049def test_parse_array_of_cftime_strings(): 1050 from cftime import DatetimeNoLeap 1051 1052 strings = np.array([["2000-01-01", "2000-01-02"], ["2000-01-03", "2000-01-04"]]) 1053 expected = np.array( 1054 [ 1055 [DatetimeNoLeap(2000, 1, 1), DatetimeNoLeap(2000, 1, 2)], 1056 [DatetimeNoLeap(2000, 1, 3), DatetimeNoLeap(2000, 1, 4)], 1057 ] 1058 ) 1059 1060 result = _parse_array_of_cftime_strings(strings, DatetimeNoLeap) 1061 np.testing.assert_array_equal(result, expected) 1062 1063 # Test scalar array case 1064 strings = np.array("2000-01-01") 1065 expected = np.array(DatetimeNoLeap(2000, 1, 1)) 1066 result = _parse_array_of_cftime_strings(strings, DatetimeNoLeap) 1067 np.testing.assert_array_equal(result, expected) 1068 1069 1070@requires_cftime 1071@pytest.mark.parametrize("calendar", _ALL_CALENDARS) 1072def test_strftime_of_cftime_array(calendar): 1073 date_format = "%Y%m%d%H%M" 1074 cf_values = xr.cftime_range("2000", periods=5, calendar=calendar) 1075 dt_values = pd.date_range("2000", periods=5) 1076 expected = pd.Index(dt_values.strftime(date_format)) 1077 result = cf_values.strftime(date_format) 1078 assert result.equals(expected) 1079 1080 1081@requires_cftime 1082@pytest.mark.parametrize("calendar", _ALL_CALENDARS) 1083@pytest.mark.parametrize("unsafe", [False, True]) 1084def test_to_datetimeindex(calendar, unsafe): 1085 index = xr.cftime_range("2000", periods=5, calendar=calendar) 1086 expected = pd.date_range("2000", periods=5) 1087 1088 if calendar in _NON_STANDARD_CALENDARS and not unsafe: 1089 with pytest.warns(RuntimeWarning, match="non-standard"): 1090 result = index.to_datetimeindex() 1091 else: 1092 result = index.to_datetimeindex(unsafe=unsafe) 1093 1094 assert result.equals(expected) 1095 np.testing.assert_array_equal(result, expected) 1096 assert isinstance(result, pd.DatetimeIndex) 1097 1098 1099@requires_cftime 1100@pytest.mark.parametrize("calendar", _ALL_CALENDARS) 1101def test_to_datetimeindex_out_of_range(calendar): 1102 index = xr.cftime_range("0001", periods=5, calendar=calendar) 1103 with pytest.raises(ValueError, match="0001"): 1104 index.to_datetimeindex() 1105 1106 1107@requires_cftime 1108@pytest.mark.parametrize("calendar", ["all_leap", "360_day"]) 1109def test_to_datetimeindex_feb_29(calendar): 1110 index = xr.cftime_range("2001-02-28", periods=2, calendar=calendar) 1111 with pytest.raises(ValueError, match="29"): 1112 index.to_datetimeindex() 1113 1114 1115@requires_cftime 1116@pytest.mark.xfail(reason="https://github.com/pandas-dev/pandas/issues/24263") 1117def test_multiindex(): 1118 index = xr.cftime_range("2001-01-01", periods=100, calendar="360_day") 1119 mindex = pd.MultiIndex.from_arrays([index]) 1120 assert mindex.get_loc("2001-01") == slice(0, 30) 1121 1122 1123@requires_cftime 1124@pytest.mark.parametrize("freq", ["3663S", "33T", "2H"]) 1125@pytest.mark.parametrize("method", ["floor", "ceil", "round"]) 1126def test_rounding_methods_against_datetimeindex(freq, method): 1127 expected = pd.date_range("2000-01-02T01:03:51", periods=10, freq="1777S") 1128 expected = getattr(expected, method)(freq) 1129 result = xr.cftime_range("2000-01-02T01:03:51", periods=10, freq="1777S") 1130 result = getattr(result, method)(freq).to_datetimeindex() 1131 assert result.equals(expected) 1132 1133 1134@requires_cftime 1135@pytest.mark.parametrize("method", ["floor", "ceil", "round"]) 1136def test_rounding_methods_invalid_freq(method): 1137 index = xr.cftime_range("2000-01-02T01:03:51", periods=10, freq="1777S") 1138 with pytest.raises(ValueError, match="fixed"): 1139 getattr(index, method)("MS") 1140 1141 1142@pytest.fixture 1143def rounding_index(date_type): 1144 return xr.CFTimeIndex( 1145 [ 1146 date_type(1, 1, 1, 1, 59, 59, 999512), 1147 date_type(1, 1, 1, 3, 0, 1, 500001), 1148 date_type(1, 1, 1, 7, 0, 6, 499999), 1149 ] 1150 ) 1151 1152 1153@requires_cftime 1154def test_ceil(rounding_index, date_type): 1155 result = rounding_index.ceil("S") 1156 expected = xr.CFTimeIndex( 1157 [ 1158 date_type(1, 1, 1, 2, 0, 0, 0), 1159 date_type(1, 1, 1, 3, 0, 2, 0), 1160 date_type(1, 1, 1, 7, 0, 7, 0), 1161 ] 1162 ) 1163 assert result.equals(expected) 1164 1165 1166@requires_cftime 1167def test_floor(rounding_index, date_type): 1168 result = rounding_index.floor("S") 1169 expected = xr.CFTimeIndex( 1170 [ 1171 date_type(1, 1, 1, 1, 59, 59, 0), 1172 date_type(1, 1, 1, 3, 0, 1, 0), 1173 date_type(1, 1, 1, 7, 0, 6, 0), 1174 ] 1175 ) 1176 assert result.equals(expected) 1177 1178 1179@requires_cftime 1180def test_round(rounding_index, date_type): 1181 result = rounding_index.round("S") 1182 expected = xr.CFTimeIndex( 1183 [ 1184 date_type(1, 1, 1, 2, 0, 0, 0), 1185 date_type(1, 1, 1, 3, 0, 2, 0), 1186 date_type(1, 1, 1, 7, 0, 6, 0), 1187 ] 1188 ) 1189 assert result.equals(expected) 1190 1191 1192@requires_cftime 1193def test_asi8(date_type): 1194 index = xr.CFTimeIndex([date_type(1970, 1, 1), date_type(1970, 1, 2)]) 1195 result = index.asi8 1196 expected = 1000000 * 86400 * np.array([0, 1]) 1197 np.testing.assert_array_equal(result, expected) 1198 1199 1200@requires_cftime 1201def test_asi8_distant_date(): 1202 """Test that asi8 conversion is truly exact.""" 1203 import cftime 1204 1205 date_type = cftime.DatetimeProlepticGregorian 1206 index = xr.CFTimeIndex([date_type(10731, 4, 22, 3, 25, 45, 123456)]) 1207 result = index.asi8 1208 expected = np.array([1000000 * 86400 * 400 * 8000 + 12345 * 1000000 + 123456]) 1209 np.testing.assert_array_equal(result, expected) 1210 1211 1212@requires_cftime 1213def test_infer_freq_valid_types(): 1214 cf_indx = xr.cftime_range("2000-01-01", periods=3, freq="D") 1215 assert xr.infer_freq(cf_indx) == "D" 1216 assert xr.infer_freq(xr.DataArray(cf_indx)) == "D" 1217 1218 pd_indx = pd.date_range("2000-01-01", periods=3, freq="D") 1219 assert xr.infer_freq(pd_indx) == "D" 1220 assert xr.infer_freq(xr.DataArray(pd_indx)) == "D" 1221 1222 pd_td_indx = pd.timedelta_range(start="1D", periods=3, freq="D") 1223 assert xr.infer_freq(pd_td_indx) == "D" 1224 assert xr.infer_freq(xr.DataArray(pd_td_indx)) == "D" 1225 1226 1227@requires_cftime 1228def test_infer_freq_invalid_inputs(): 1229 # Non-datetime DataArray 1230 with pytest.raises(ValueError, match="must contain datetime-like objects"): 1231 xr.infer_freq(xr.DataArray([0, 1, 2])) 1232 1233 indx = xr.cftime_range("1990-02-03", periods=4, freq="MS") 1234 # 2D DataArray 1235 with pytest.raises(ValueError, match="must be 1D"): 1236 xr.infer_freq(xr.DataArray([indx, indx])) 1237 1238 # CFTimeIndex too short 1239 with pytest.raises(ValueError, match="Need at least 3 dates to infer frequency"): 1240 xr.infer_freq(indx[:2]) 1241 1242 # Non-monotonic input 1243 assert xr.infer_freq(indx[np.array([0, 2, 1, 3])]) is None 1244 1245 # Non-unique input 1246 assert xr.infer_freq(indx[np.array([0, 1, 1, 2])]) is None 1247 1248 # No unique frequency (here 1st step is MS, second is 2MS) 1249 assert xr.infer_freq(indx[np.array([0, 1, 3])]) is None 1250 1251 # Same, but for QS 1252 indx = xr.cftime_range("1990-02-03", periods=4, freq="QS") 1253 assert xr.infer_freq(indx[np.array([0, 1, 3])]) is None 1254 1255 1256@requires_cftime 1257@pytest.mark.parametrize( 1258 "freq", 1259 [ 1260 "300AS-JAN", 1261 "A-DEC", 1262 "AS-JUL", 1263 "2AS-FEB", 1264 "Q-NOV", 1265 "3QS-DEC", 1266 "MS", 1267 "4M", 1268 "7D", 1269 "D", 1270 "30H", 1271 "5T", 1272 "40S", 1273 ], 1274) 1275@pytest.mark.parametrize("calendar", _CFTIME_CALENDARS) 1276def test_infer_freq(freq, calendar): 1277 indx = xr.cftime_range("2000-01-01", periods=3, freq=freq, calendar=calendar) 1278 out = xr.infer_freq(indx) 1279 assert out == freq 1280