diff --git a/src/earthkit/time/cli/sequence.py b/src/earthkit/time/cli/sequence.py index 3857130..0052725 100644 --- a/src/earthkit/time/cli/sequence.py +++ b/src/earthkit/time/cli/sequence.py @@ -23,6 +23,11 @@ def seq_prev_action(parser: argparse.ArgumentParser, args: argparse.Namespace): print(format_date(seq.previous(args.date, strict=(not args.inclusive)))) +def seq_nearest_action(parser: argparse.ArgumentParser, args: argparse.Namespace): + seq = create_sequence(parser, args) + print(format_date(seq.nearest(args.date, resolve=args.resolve))) + + def seq_range_action(parser: argparse.ArgumentParser, args: argparse.Namespace): seq = create_sequence(parser, args) print( @@ -87,6 +92,23 @@ def get_parser() -> argparse.ArgumentParser: help="if the given date is in the sequence, return it", ) + nearest_action = parser.add_action( + "nearest", + seq_nearest_action, + help="compute the nearest date in the given sequence", + description="Compute the nearest date in the given sequence", + epilog=SEQ_EPILOG, + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + nearest_action.add_argument("date", type=parse_date, help="reference date") + add_sequence_args(nearest_action) + nearest_action.add_argument( + "--resolve", + choices=("previous", "next"), + default="previous", + help="return this date in case of a tie", + ) + range_action = parser.add_action( "range", seq_range_action, diff --git a/src/earthkit/time/sequence.py b/src/earthkit/time/sequence.py index 501f70f..4bfb70b 100644 --- a/src/earthkit/time/sequence.py +++ b/src/earthkit/time/sequence.py @@ -51,6 +51,26 @@ def previous(self, reference: date, strict: bool = True) -> date: current -= oneday return current + def nearest(self, reference: date, resolve: str = "previous") -> date: + """Return the date closest to ``reference`` in the sequence. + In case this is ambiguous, ``resolve`` defines which date to use + (``"previous"`` or ``"next"``). + """ + if resolve not in ["previous", "next"]: + raise ValueError('`resolve` must be either "previous" or "next"') + before = self.previous(reference, strict=False) + after = self.next(reference, strict=False) + delta_b = reference - before + delta_a = after - reference + if delta_b < delta_a: + return before + elif delta_b > delta_a: + return after + elif resolve == "previous": + return before + else: + return after + def range( self, start: date, diff --git a/tests/cli/test_sequence_cli.py b/tests/cli/test_sequence_cli.py index 4d5b75c..6afac88 100644 --- a/tests/cli/test_sequence_cli.py +++ b/tests/cli/test_sequence_cli.py @@ -6,6 +6,7 @@ from earthkit.time.calendar import Weekday from earthkit.time.cli.sequence import ( seq_bracket_action, + seq_nearest_action, seq_next_action, seq_prev_action, seq_range_action, @@ -104,6 +105,56 @@ def test_seq_prev(args: dict, expected: str, capsys: pytest.CaptureFixture[str]) assert captured.out == expected + "\n" +@pytest.mark.parametrize( + "args, expected", + [ + pytest.param( + {"daily": True, "date": date(2006, 7, 26)}, "20060726", id="daily" + ), + pytest.param( + {"daily": True, "date": date(2017, 3, 30), "exclude": ["30", "31"]}, + "20170329", + id="daily-excludes", + ), + pytest.param( + { + "weekly": [Weekday.TUESDAY, Weekday.THURSDAY, Weekday.SATURDAY], + "date": date(2013, 10, 23), + "resolve": "previous", + }, + "20131022", + id="weekly", + ), + pytest.param( + {"monthly": [1, 15], "date": date(1995, 8, 25)}, + "19950901", + id="monthly", + ), + pytest.param( + { + "yearly": [(1, 4), (12, 25)], + "date": date(2009, 12, 30), + "resolve": "next", + }, + "20100104", + id="yearly", + ), + ], +) +def test_seq_nearest(args: dict, expected: str, capsys: pytest.CaptureFixture[str]): + parser = argparse.ArgumentParser() + args.setdefault("daily", False) + args.setdefault("weekly", None) + args.setdefault("monthly", None) + args.setdefault("yearly", None) + args.setdefault("exclude", []) + args.setdefault("resolve", "previous") + args = argparse.Namespace(**args) + seq_nearest_action(parser, args) + captured = capsys.readouterr() + assert captured.out == expected + "\n" + + @pytest.mark.parametrize( "args, expected", [ diff --git a/tests/test_sequence.py b/tests/test_sequence.py index 1f7d7df..17bb4e9 100644 --- a/tests/test_sequence.py +++ b/tests/test_sequence.py @@ -45,7 +45,7 @@ def __contains__(self, day: date) -> bool: pytest.param( DailySequence(), [(1983, 4, 28), (1983, 4, 29), (1983, 4, 30), (1983, 5, 1), (1983, 5, 2)], - None, + [], id="daily-simple", ), pytest.param( @@ -57,43 +57,43 @@ def __contains__(self, day: date) -> bool: (2002, 1, 1), (2002, 1, 2), ], - None, + [], id="daily-crossyear", ), pytest.param( DailySequence(), [(2003, 2, 26), (2003, 2, 27), (2003, 2, 28), (2003, 3, 1), (2003, 3, 2)], - None, + [], id="daily-feb28-nonleap", ), pytest.param( DailySequence(), [(2016, 2, 27), (2016, 2, 28), (2016, 2, 29), (2016, 3, 1), (2016, 3, 2)], - None, + [], id="daily-feb28-leap", ), pytest.param( DailySequence(excludes=[1]), [(1995, 4, 28), (1995, 4, 29), (1995, 4, 30), (1995, 5, 2), (1995, 5, 3)], - (1995, 5, 1), + [(1995, 5, 1)], id="daily-exclude", ), pytest.param( DailySequence(excludes=[29]), [(2000, 2, 26), (2000, 2, 27), (2000, 2, 28), (2000, 3, 1), (2000, 3, 2)], - (2000, 2, 29), + [(2000, 2, 29)], id="daily-exclude-leap", ), pytest.param( DailySequence(excludes=range(1, 29)), [(2020, 1, 30), (2020, 1, 31), (2020, 2, 29), (2020, 3, 29), (2020, 3, 30)], - (2020, 2, 20), + [(2020, 2, 20)], id="daily-exclude-almostall", ), pytest.param( WeeklySequence(2), [(1999, 3, 24), (1999, 3, 31), (1999, 4, 7), (1999, 4, 14), (1999, 4, 21)], - (1999, 4, 20), + [(1999, 4, 20)], id="weekly-simple", ), pytest.param( @@ -105,25 +105,25 @@ def __contains__(self, day: date) -> bool: (2012, 1, 13), (2012, 1, 20), ], - (2012, 1, 2), + [(2012, 1, 2)], id="weekly-crossyear", ), pytest.param( WeeklySequence([2, 5]), [(2007, 2, 24), (2007, 2, 28), (2007, 3, 3), (2007, 3, 7), (2007, 3, 10)], - (2007, 2, 25), + [(2007, 2, 25), (2007, 3, 5)], id="weekly-feb28-nonleap", ), pytest.param( WeeklySequence([MONDAY, THURSDAY]), [(2024, 2, 22), (2024, 2, 26), (2024, 2, 29), (2024, 3, 4), (2024, 3, 7)], - (2024, 2, 28), + [(2024, 2, 28), (2024, 3, 2)], id="weekly-feb28-leap", ), pytest.param( MonthlySequence(15), [(1989, 3, 15), (1989, 4, 15), (1989, 5, 15), (1989, 6, 15), (1989, 7, 15)], - (1989, 5, 19), + [(1989, 4, 30), (1989, 5, 19)], id="monthly-simple", ), pytest.param( @@ -135,19 +135,19 @@ def __contains__(self, day: date) -> bool: (2015, 1, 7), (2015, 1, 21), ], - (2014, 12, 31), + [(2014, 12, 14), (2014, 12, 31)], id="monthly-crossyear", ), pytest.param( MonthlySequence(range(1, 32, 7)), [(2009, 2, 15), (2009, 2, 22), (2009, 3, 1), (2009, 3, 8), (2009, 3, 15)], - (2009, 3, 14), + [(2009, 3, 14)], id="monthly-feb28-nonleap", ), pytest.param( MonthlySequence([28, 29]), [(1992, 1, 29), (1992, 2, 28), (1992, 2, 29), (1992, 3, 28), (1992, 3, 29)], - (1992, 2, 18), + [(1992, 2, 18), (1992, 3, 14)], id="monthly-feb28-leap", ), pytest.param( @@ -159,13 +159,13 @@ def __contains__(self, day: date) -> bool: (1987, 12, 11), (1987, 12, 22), ], - (1987, 11, 11), + [(1987, 11, 11)], id="monthly-exclude", ), pytest.param( MonthlySequence([27, 29, 31], excludes=[(2, 29)]), [(2008, 1, 31), (2008, 2, 27), (2008, 3, 27), (2008, 3, 29), (2008, 3, 31)], - (2008, 2, 29), + [(2008, 2, 29), (2008, 3, 28)], id="monthly-exclude-leap", ), pytest.param( @@ -177,13 +177,13 @@ def __contains__(self, day: date) -> bool: (2022, 12, 31), (2023, 10, 31), ], - (2022, 5, 9), + [(2022, 5, 9), (2022, 6, 1)], id="monthly-exclude-almostall", ), pytest.param( YearlySequence((1, 1)), [(1999, 1, 1), (2000, 1, 1), (2001, 1, 1), (2002, 1, 1), (2003, 1, 1)], - (2002, 2, 2), + [(2000, 7, 2), (2002, 2, 2)], id="yearly-simple", ), pytest.param( @@ -195,43 +195,43 @@ def __contains__(self, day: date) -> bool: (2018, 4, 2), (2018, 7, 2), ], - (2018, 6, 18), + [(2017, 11, 16), (2018, 6, 18)], id="yearly-crossyear", ), pytest.param( YearlySequence([(2, 28), (2, 29), (3, 1)]), [(1994, 2, 28), (1994, 3, 1), (1995, 2, 28), (1995, 3, 1), (1996, 2, 28)], - (1995, 2, 22), + [(1994, 8, 30), (1995, 2, 22)], id="yearly-feb28-nonleap", ), pytest.param( YearlySequence([(i, 29) for i in range(1, 13)]), [(2008, 1, 29), (2008, 2, 29), (2008, 3, 29), (2008, 4, 29), (2008, 5, 29)], - (2008, 2, 28), + [(2008, 2, 28), (2008, 5, 14)], id="yearly-feb28-leap", ), pytest.param( YearlySequence((2, 29)), [(2000, 2, 29), (2004, 2, 29), (2008, 2, 29), (2012, 2, 29), (2016, 2, 29)], - (2003, 2, 28), + [(2003, 2, 28)], id="yearly-leaponly", ), pytest.param( YearlySequence([(7, 13)], excludes=[date(2023, 7, 13)]), [(2020, 7, 13), (2021, 7, 13), (2022, 7, 13), (2024, 7, 13), (2025, 7, 13)], - (2023, 7, 13), + [(2023, 7, 13)], id="yearly-exclude", ), pytest.param( YearlySequence([(1, 31), (2, 28), (2, 29)], excludes=ExcludeLeapFeb28()), [(2007, 1, 31), (2007, 2, 28), (2008, 1, 31), (2008, 2, 29), (2009, 1, 31)], - (2008, 3, 10), + [(2007, 2, 14), (2008, 3, 10)], id="yearly-exclude-leap", ), pytest.param( YearlySequence([(3, 31)], excludes=ExcludeNonTenYears()), # FIXME [(1990, 3, 31), (2000, 3, 31), (2010, 3, 31), (2020, 3, 31), (2030, 3, 31)], - (2001, 3, 31), + [(2001, 3, 31), (2005, 3, 31)], id="yearly-exclude-almostall", ), ], @@ -239,11 +239,9 @@ def __contains__(self, day: date) -> bool: def test_sequence( seq: Sequence, ymds: List[Tuple[int, int, int]], - outside: Optional[Tuple[int, int, int]], + outside: List[Tuple[int, int, int]], ): dates = [date(y, m, d) for y, m, d in ymds] - if outside is not None: - outside = date(*outside) for d in dates: assert d in seq @@ -257,6 +255,9 @@ def test_sequence( assert seq.next(cur, False) == cur assert seq.previous(cur) == prev assert seq.previous(cur, False) == cur + assert seq.nearest(cur) == cur + assert seq.nearest(cur, resolve="previous") == cur + assert seq.nearest(cur, resolve="next") == cur assert list(seq.range(dates[0], dates[-1])) == dates assert list(seq.range(dates[0], dates[-1], include_start=False)) == dates[1:] @@ -272,28 +273,41 @@ def test_sequence( assert list(seq.bracket(dates[2], 1, strict=False)) == dates[1:4] assert list(seq.bracket(dates[2], (2, 1), strict=False)) == dates[:4] - if outside is not None: - out_i = bisect.bisect_left(dates, outside) + for out_tup in outside: + out_date = date(*out_tup) + out_i = bisect.bisect_left(dates, out_date) before = out_i after = len(dates) - out_i - assert outside not in seq - assert seq.next(outside) == dates[out_i] - assert seq.previous(outside) == dates[out_i - 1] + assert out_date not in seq + assert seq.next(out_date) == dates[out_i] + assert seq.previous(out_date) == dates[out_i - 1] - assert list(seq.range(outside, dates[-1])) == dates[out_i:] - assert list(seq.range(outside, dates[-1], include_start=False)) == dates[out_i:] - assert list(seq.range(dates[0], outside)) == dates[:out_i] - assert list(seq.range(dates[0], outside, include_end=False)) == dates[:out_i] + db = (out_date - dates[out_i - 1]).days + da = (dates[out_i] - out_date).days + if db != da: + exp_nearest = dates[out_i - 1] if db < da else dates[out_i] + assert seq.nearest(out_date) == exp_nearest + else: + assert seq.nearest(out_date) == dates[out_i - 1] + assert seq.nearest(out_date, resolve="previous") == dates[out_i - 1] + assert seq.nearest(out_date, resolve="next") == dates[out_i] + + assert list(seq.range(out_date, dates[-1])) == dates[out_i:] + assert ( + list(seq.range(out_date, dates[-1], include_start=False)) == dates[out_i:] + ) + assert list(seq.range(dates[0], out_date)) == dates[:out_i] + assert list(seq.range(dates[0], out_date, include_end=False)) == dates[:out_i] - assert list(seq.bracket(outside)) == [dates[out_i - 1], dates[out_i]] - assert list(seq.bracket(outside, (before, after))) == dates + assert list(seq.bracket(out_date)) == [dates[out_i - 1], dates[out_i]] + assert list(seq.bracket(out_date, (before, after))) == dates assert ( - list(seq.bracket(outside, (min(2, before), after))) + list(seq.bracket(out_date, (min(2, before), after))) == dates[max(0, out_i - 2) :] ) assert ( - list(seq.bracket(outside, (before, min(2, after)))) + list(seq.bracket(out_date, (before, min(2, after)))) == dates[: min(len(dates), out_i + 2)] )