Handle ambiguity in --min-date and --max-date

This fixes a bug where incomplete dates such as --max-date 2018 would
not be inclusive, since that had previously resolved to --max-date
2018-01-01.

numeric_date() and get_numerical_date_from_value() are two separate
functions that serve similar purposes: converting some date value to a
numeric date. The latter has had more recent developments, and crucially
will return a date range for ambiguous dates. This is desired for
min/max date filters, so I've switched to it.

Author: victorlin

augur/dates/__init__.py

@@ -6,7 +6,7 @@
 import re
 from functools import cache
 from treetime.utils import numeric_date as tt_numeric_date, datetime_from_numeric
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, Literal, Optional, Tuple, Union
 from augur.errors import AugurError
 from .errors import InvalidDate
 
@@ -75,17 +75,49 @@ def numeric_date(date):
 
     raise InvalidDate(date, f"""Ensure it is in one of the supported formats:\n{SUPPORTED_DATE_HELP_TEXT}""")
 
-def numeric_date_type(date):
-    """Wraps numeric_date() for argparse usage.
+def numeric_date_type_min(date) -> float:
+    """Wraps numeric_date() for argparse usage, taking the minimum value if resolved to a range.
 
     This raises an ArgumentTypeError from InvalidDateFormat exceptions, otherwise the custom exception message won't be shown in console output due to:
     https://github.com/python/cpython/blob/5c4d1f6e0e192653560ae2941a6677fbf4fbd1f2/Lib/argparse.py#L2503-L2513
+
+    >>> round(numeric_date_type_min("2018"), 3)
+    2018.001
     """
     try:
-        return numeric_date(date)
+        return get_single_numeric_date(date, fmt="%Y-%m-%d", min_or_max="min")
     except InvalidDate as error:
         raise argparse.ArgumentTypeError(str(error)) from error
 
+def numeric_date_type_max(date) -> float:
+    """Wraps numeric_date() for argparse usage, taking the maximum value if resolved to a range.
+
+    This raises an ArgumentTypeError from InvalidDateFormat exceptions, otherwise the custom exception message won't be shown in console output due to:
+    https://github.com/python/cpython/blob/5c4d1f6e0e192653560ae2941a6677fbf4fbd1f2/Lib/argparse.py#L2503-L2513
+
+    >>> round(numeric_date_type_max("2018"), 3)
+    2018.999
+    """
+    try:
+        return get_single_numeric_date(date, fmt="%Y-%m-%d", min_or_max="max")
+    except InvalidDate as error:
+        raise argparse.ArgumentTypeError(str(error)) from error
+
+def get_single_numeric_date(value, fmt, min_or_max: Literal["min", "max"]) -> float:
+    numeric_date = get_numerical_date_from_value(value, fmt)
+
+    if isinstance(numeric_date, float):
+        return numeric_date
+
+    if isinstance(numeric_date, tuple):
+        if min_or_max == "min":
+            return numeric_date[0]
+        if min_or_max == "max":
+            return numeric_date[1]
+
+    raise InvalidDate(value, f"""Ensure it is in one of the supported formats:\n{SUPPORTED_DATE_HELP_TEXT}""")
+
+
 def is_date_ambiguous(date, ambiguous_by):
     """
     Returns whether a given date string in the format of YYYY-MM-DD is ambiguous by a given part of the date (e.g., day, month, year, or any parts).

augur/filter/__init__.py

@@ -12,7 +12,7 @@
 program vcftools must be available on PATH.
 """
 from augur.argparse_ import ExtendOverwriteDefault, SKIP_AUTO_DEFAULT_IN_HELP
-from augur.dates import numeric_date_type
+from augur.dates import numeric_date_type_min, numeric_date_type_max
 from augur.filter.arguments import descriptions
 from augur.filter.io import column_type_pair
 from augur.io.metadata import DEFAULT_DELIMITERS, DEFAULT_ID_COLUMNS
@@ -39,8 +39,8 @@ def register_arguments(parser):
 
     metadata_filter_group.add_argument('--query', help=descriptions['query'])
     metadata_filter_group.add_argument('--query-columns', type=column_type_pair, nargs="+", action=ExtendOverwriteDefault, help=descriptions['query_columns'])
-    metadata_filter_group.add_argument('--min-date', type=numeric_date_type, help=descriptions['min_date'])
-    metadata_filter_group.add_argument('--max-date', type=numeric_date_type, help=descriptions['max_date'])
+    metadata_filter_group.add_argument('--min-date', type=numeric_date_type_min, help=descriptions['min_date'])
+    metadata_filter_group.add_argument('--max-date', type=numeric_date_type_max, help=descriptions['max_date'])
     metadata_filter_group.add_argument('--exclude-ambiguous-dates-by', choices=['any', 'day', 'month', 'year'], help=descriptions['exclude_ambiguous_dates_by'])
     metadata_filter_group.add_argument('--exclude', type=str, nargs="+", action=ExtendOverwriteDefault, help=descriptions['exclude'])
     metadata_filter_group.add_argument('--exclude-where', nargs='+', action=ExtendOverwriteDefault, help=descriptions['exclude_where'])

augur/frequencies.py

@@ -10,7 +10,7 @@
 from .errors import AugurError
 from .frequency_estimators import get_pivots, alignment_frequencies, tree_frequencies
 from .frequency_estimators import AlignmentKdeFrequencies, TreeKdeFrequencies, TreeKdeFrequenciesError
-from .dates import numeric_date_type, SUPPORTED_DATE_HELP_TEXT, get_numerical_dates
+from .dates import numeric_date_type_min, numeric_date_type_max, SUPPORTED_DATE_HELP_TEXT, get_numerical_dates
 from .io.file import open_file
 from .io.metadata import DEFAULT_DELIMITERS, DEFAULT_ID_COLUMNS, METADATA_DATE_COLUMN, InvalidDelimiter, Metadata, read_metadata
 from .utils import write_json
@@ -37,9 +37,9 @@ def register_parser(parent_subparsers):
                         help="number of units between pivots")
     parser.add_argument("--pivot-interval-units", type=str, default="months", choices=['months', 'weeks'],
                         help="space pivots by months (default) or by weeks")
-    parser.add_argument('--min-date', type=numeric_date_type,
+    parser.add_argument('--min-date', type=numeric_date_type_min,
                         help=f"date to begin frequencies calculations; may be specified as: {SUPPORTED_DATE_HELP_TEXT}")
-    parser.add_argument('--max-date', type=numeric_date_type,
+    parser.add_argument('--max-date', type=numeric_date_type_max,
                         help=f"date to end frequencies calculations; may be specified as: {SUPPORTED_DATE_HELP_TEXT}")
 
     # Tree-specific arguments

tests/functional/filter/cram/filter-max-date.t

@@ -6,16 +6,16 @@ Create metadata TSV file for testing.
 
   $ cat >metadata.tsv <<~~
   > strain	date
-  > SEQ_1	2020-03-XX
-  > SEQ_2	2020-03-01
-  > SEQ_3	2020-03-02
+  > SEQ_1	2019-XX-XX
+  > SEQ_2	2019-12-31
+  > SEQ_3	2020-01-01
   > ~~
 
-Test that --max-date is inclusive.
+Test that --max-date is inclusive even with ambiguity.
 
   $ ${AUGUR} filter \
   >  --metadata metadata.tsv \
-  >  --max-date 2020-03-01 \
+  >  --max-date 2019 \
   >  --output-strains filtered_strains.txt 2>/dev/null
   $ sort filtered_strains.txt
   SEQ_1