Skip to content

Commit

Permalink
Merge pull request #3 from ecmwf/feature/sequence-nearest
Browse files Browse the repository at this point in the history
Add "nearest" method to Sequence
  • Loading branch information
oiffrig authored Jul 19, 2024
2 parents c4fb9c7 + ab18ef0 commit 8401135
Show file tree
Hide file tree
Showing 4 changed files with 149 additions and 42 deletions.
22 changes: 22 additions & 0 deletions src/earthkit/time/cli/sequence.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down
20 changes: 20 additions & 0 deletions src/earthkit/time/sequence.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
51 changes: 51 additions & 0 deletions tests/cli/test_sequence_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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",
[
Expand Down
Loading

0 comments on commit 8401135

Please sign in to comment.