diff --git a/conservancy_beancount/cliutil.py b/conservancy_beancount/cliutil.py index 3e7138e..73268d9 100644 --- a/conservancy_beancount/cliutil.py +++ b/conservancy_beancount/cliutil.py @@ -46,6 +46,7 @@ from typing import ( IO, Iterable, Iterator, + List, NamedTuple, NoReturn, Optional, @@ -195,6 +196,42 @@ class ExitCode(enum.IntEnum): BeancountErrors = 63 +class ExtendAction(argparse.Action): + """argparse action to let a user build a list from a string + + This is a fancier version of argparse's built-in ``action='append'``. + The user's input is turned into a list of strings, split by a regexp + pattern you provide. Typical usage looks like:: + + parser = argparse.ArgumentParser() + parser.add_argument( + '--option', ..., + action=ExtendAction, + const=regexp_pattern, # default is r'\s*,\s*' + ..., + ) + """ + DEFAULT_PATTERN = r'\s*,\s*' + + def __call__(self, + parser: argparse.ArgumentParser, + namespace: argparse.Namespace, + values: Union[Sequence[Any], str, None]=None, + option_string: Optional[str]=None, + ) -> None: + pattern: str = self.const or self.DEFAULT_PATTERN + value: Optional[List[str]] = getattr(namespace, self.dest, None) + if value is None: + value = [] + setattr(namespace, self.dest, value) + if values is None: + values = [] + elif isinstance(values, str): + values = [values] + for s in values: + value.extend(re.split(pattern, s)) + + class InfoAction(argparse.Action): def __call__(self, parser: argparse.ArgumentParser, diff --git a/conservancy_beancount/reconcile/statement.py b/conservancy_beancount/reconcile/statement.py index 4c08643..b2039df 100644 --- a/conservancy_beancount/reconcile/statement.py +++ b/conservancy_beancount/reconcile/statement.py @@ -338,10 +338,10 @@ The default is one month after the start date. '--account', '-a', dest='accounts', metavar='ACCOUNT', - action='append', - help="""Reconcile this account. You can specify this option -multiple times. You can specify a part of the account hierarchy, or an account -classification from metadata. Default adapts to your search criteria. + action=cliutil.ExtendAction, + help="""Reconcile this account. You can specify multiple +comma-separated accounts or classifications, and/or specify this option +multiple times. Default adapts to your search criteria. """) parser.add_argument( '--id-metadata-key', '-i', diff --git a/conservancy_beancount/reports/ledger.py b/conservancy_beancount/reports/ledger.py index 9af320c..b90444d 100644 --- a/conservancy_beancount/reports/ledger.py +++ b/conservancy_beancount/reports/ledger.py @@ -739,8 +739,9 @@ date was also not specified. '--account', '-a', dest='accounts', metavar='ACCOUNT', - action='append', - help="""Show this account in the report. You can specify this option + action=cliutil.ExtendAction, + help="""Show this account in the report. You can specify multiple +comma-separated accounts or classifications, and/or specify this option multiple times. You can specify a part of the account hierarchy, or an account classification from metadata. If not specified, the default set adapts to your search criteria. @@ -749,19 +750,21 @@ search criteria. parser.add_argument( '--show-totals', '-S', metavar='ACCOUNT', - action='append', + action=cliutil.ExtendAction, help="""When entries for this account appear in the report, include -account balance(s) as well. You can specify this option multiple times. Pass in -a part of the account hierarchy. The default is all accounts. +account balance(s) as well. You can specify multiple comma-separated parts of +the account hierarchy, and/or specify this option multiple times. +The default is all accounts. """) parser.add_argument( '--add-totals', '-T', metavar='ACCOUNT', - action='append', + action=cliutil.ExtendAction, help="""When an account could be included in the report but does not have any entries in the date range, include a header and account balance(s) for -it. You can specify this option multiple times. Pass in a part of the account -hierarchy. The default set adapts to your search criteria. +it. You can specify multiple comma-separated parts of the account hierarchy, +and/or specify this option multiple times. The default set adapts to your +search criteria. """) parser.add_argument( '--sheet-size', '--size', diff --git a/tests/test_cliutil.py b/tests/test_cliutil.py index 510d615..b0425ed 100644 --- a/tests/test_cliutil.py +++ b/tests/test_cliutil.py @@ -298,3 +298,36 @@ def test_enum_arg_choices_str_defaults(arg_enum): def test_enum_arg_choices_str_args(arg_enum): sep = '/' assert arg_enum.choices_str(sep, '{}') == sep.join(c.value for c in ArgChoices) + +@pytest.mark.parametrize('values,sep', testutil.combine_values( + [['foo'], ['bar', 'baz'], ['qu', 'quu', 'quux']], + [',', ', ', ' ,', ' , '], +)) +def test_extend_action_once(values, sep): + action = cliutil.ExtendAction(['-t'], 'result') + args = argparse.Namespace() + action(None, args, sep.join(values), '-t') + assert args.result == values + +def test_extend_action_multi(): + action = cliutil.ExtendAction(['-t'], 'result') + args = argparse.Namespace() + action(None, args, 'foo,bar', '-t') + action(None, args, 'baz, quux', '-t') + assert args.result == ['foo', 'bar', 'baz', 'quux'] + +def test_extend_action_from_default(): + action = cliutil.ExtendAction(['-t'], 'result') + args = argparse.Namespace(result=['foo']) + action(None, args, 'bar , baz', '-t') + assert args.result == ['foo', 'bar', 'baz'] + +@pytest.mark.parametrize('pattern,expected', [ + (',', ['foo', ' bar']), + (r'\s+', ['foo,', 'bar']), +]) +def test_extend_action_custom_pattern(pattern, expected): + action = cliutil.ExtendAction(['-t'], 'result', const=pattern) + args = argparse.Namespace() + action(None, args, 'foo, bar', '-t') + assert args.result == expected