cliutil: Add year_or_date_arg() function.

conservancy_beancount/cliutil.py

@@ -235,6 +235,23 @@ def add_version_argument(parser: argparse.ArgumentParser) -> argparse.Action:
 def date_arg(arg: str) -> datetime.date:
     return datetime.datetime.strptime(arg, '%Y-%m-%d').date()
 
+def year_or_date_arg(arg: str) -> Union[int, datetime.date]:
+    """Get either a date or a year (int) from an argument string
+
+    This is a useful argument type for arguments that will be passed into
+    Books loader methods which can accept either a fiscal year or a full date.
+    """
+    try:
+        year = int(arg, 10)
+    except ValueError:
+        ok = False
+    else:
+        ok = datetime.MINYEAR <= year <= datetime.MAXYEAR
+    if ok:
+        return year
+    else:
+        return date_arg(arg)
+
 def make_entry_point(mod_name: str, prog_name: str=sys.argv[0]) -> Callable[[], int]:
     """Create an entry_point function for a tool
 

tests/test_cliutil.py

@@ -81,9 +81,8 @@ def test_bytes_output_stream(path):
     (2020, 12, 31),
 ])
 def test_date_arg_valid(year, month, day):
-    arg = f'{year}-{month}-{day}'
     expected = datetime.date(year, month, day)
-    assert cliutil.date_arg(arg) == expected
+    assert cliutil.date_arg(expected.isoformat()) == expected
 
 @pytest.mark.parametrize('arg', [
     '2000',
@@ -95,6 +94,33 @@ def test_date_arg_invalid(arg):
     with pytest.raises(ValueError):
         cliutil.date_arg(arg)
 
+@pytest.mark.parametrize('year', [
+    1990,
+    2000,
+    2009,
+])
+def test_year_or_date_arg_year(year):
+    assert cliutil.year_or_date_arg(str(year)) == year
+
+@pytest.mark.parametrize('year,month,day', [
+    (2000, 1, 1),
+    (2016, 2, 29),
+    (2020, 12, 31),
+])
+def test_year_or_date_arg_date(year, month, day):
+    expected = datetime.date(year, month, day)
+    assert cliutil.year_or_date_arg(expected.isoformat()) == expected
+
+@pytest.mark.parametrize('arg', [
+    '-1',
+    str(sys.maxsize),
+    'MMDVIII',
+    '2019-02-29',
+])
+def test_year_or_date_arg_invalid(arg):
+    with pytest.raises(ValueError):
+        cliutil.year_or_date_arg(arg)
+
 @pytest.mark.parametrize('func_name', [
     'bytes_output',
     'text_output',