Merge pull request #3 from ecmwf/feature/sequence-nearest

Add "nearest" method to Sequence

Author: oiffrig

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,

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,

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",
     [