diff --git a/conservancy_beancount/reports/query.py b/conservancy_beancount/reports/query.py new file mode 100644 index 0000000000000000000000000000000000000000..302423f5c9335d45da502cbfb9c289912d0fca8c --- /dev/null +++ b/conservancy_beancount/reports/query.py @@ -0,0 +1,236 @@ +"""query.py - Report arbitrary queries with advanced loading and formatting""" +# Copyright © 2021 Brett Smith +# License: AGPLv3-or-later WITH Beancount-Plugin-Additional-Permission-1.0 +# +# Full copyright and licensing details can be found at toplevel file +# LICENSE.txt in the repository. + +import argparse +import collections +import datetime +import enum +import functools +import locale +import logging +import re +import sys + +from typing import ( + cast, + Callable, + Dict, + Iterable, + Iterator, + Mapping, + Optional, + Sequence, + TextIO, + Tuple, + Union, +) +from ..beancount_types import ( + MetaValue, + Posting, + Transaction, +) + +from decimal import Decimal +from pathlib import Path + +import beancount.query.shell as bc_query +import beancount.query.query_parser as bc_query_parser + +from . import core +from . import rewrite +from .. import books +from .. import cliutil +from .. import config as configmod +from .. import data + +PROGNAME = 'query-report' +QUERY_PARSER = bc_query_parser.Parser() +logger = logging.getLogger('conservancy_beancount.reports.query') + +class BooksLoader: + """Closure to load books with a zero-argument callable + + This matches the load interface that BQLShell expects. + """ + def __init__( + self, + books_loader: Optional[books.Loader], + start_date: Optional[datetime.date]=None, + stop_date: Optional[datetime.date]=None, + rewrite_rules: Sequence[rewrite.RewriteRuleset]=(), + ) -> None: + self.books_loader = books_loader + self.start_date = start_date + self.stop_date = stop_date + self.rewrite_rules = rewrite_rules + + def __call__(self) -> books.LoadResult: + result = books.Loader.dispatch(self.books_loader, self.start_date, self.stop_date) + for index, entry in enumerate(result.entries): + # entry might not be a Transaction; we catch that later. + # The type ignores are because the underlying Beancount type isn't + # type-checkable. + postings = data.Posting.from_txn(entry) # type:ignore[arg-type] + for ruleset in self.rewrite_rules: + postings = ruleset.rewrite(postings) + try: + result.entries[index] = entry._replace(postings=list(postings)) # type:ignore[call-arg] + except AttributeError: + pass + return result + + +class BQLShell(bc_query.BQLShell): + pass + + +class JoinOperator(enum.Enum): + AND = 'AND' + OR = 'OR' + + def join(self, parts: Iterable[str]) -> str: + return f' {self.value} '.join(parts) + + +class ReportFormat(enum.Enum): + TEXT = 'text' + TXT = TEXT + CSV = 'csv' + # ODS = 'ods' + + +def _date_condition( + date: Union[int, datetime.date], + year_to_date: Callable[[int], datetime.date], + op: str, +) -> str: + if isinstance(date, int): + date = year_to_date(date) + return f'date {op} {date.isoformat()}' + +def build_query( + args: argparse.Namespace, + fy: books.FiscalYear, + in_file: Optional[TextIO]=None, +) -> Optional[str]: + if not args.query: + args.query = [] if in_file is None else [line[:-1] for line in in_file] + if not any(re.search(r'\S', s) for s in args.query): + return None + plain_query = ' '.join(args.query) + try: + QUERY_PARSER.parse(plain_query) + except bc_query_parser.ParseError: + conds = [f'({args.join.join(args.query)})'] + if args.start_date is not None: + conds.append(_date_condition(args.start_date, fy.first_date, '>=')) + if args.stop_date is not None: + conds.append(_date_condition(args.stop_date, fy.next_fy_date, '<')) + return f'SELECT * WHERE {" AND ".join(conds)}' + else: + return plain_query + +def parse_arguments(arglist: Optional[Sequence[str]]=None) -> argparse.Namespace: + parser = argparse.ArgumentParser(prog=PROGNAME) + cliutil.add_version_argument(parser) + parser.add_argument( + '--begin', '--start', '-b', + dest='start_date', + metavar='DATE', + type=cliutil.year_or_date_arg, + help="""Begin loading entries from this fiscal year. When query-report +builds the query, it will include a condition `date >= DATE`. +""") + parser.add_argument( + '--end', '--stop', '-e', + dest='stop_date', + metavar='DATE', + type=cliutil.year_or_date_arg, + help="""End loading entries from this fiscal year. When query-report +builds the query, it will include a condition `date < DATE`. If you specify a +begin date but not an end date, the default end date will be the end of the +fiscal year of the begin date. +""") + cliutil.add_rewrite_rules_argument(parser) + format_arg = cliutil.EnumArgument(ReportFormat) + parser.add_argument( + '--report-type', '--format', '-t', '-f', + metavar='TYPE', + type=format_arg.enum_type, + help="""Format of report to generate. Choices are +{format_arg.choices_str()}. Default is guessed from your output filename +extension, or 'text' if that fails. +""") + parser.add_argument( + '--output-file', '-O', '-o', + metavar='PATH', + type=Path, + help="""Write the report to this file, or stdout when PATH is `-`. +The default is stdout for text and CSV reports, and a generated filename for +ODS reports. +""") + join_arg = cliutil.EnumArgument(JoinOperator) + parser.add_argument( + '--join', '-j', + metavar='OP', + type=join_arg.enum_type, + default=JoinOperator.AND, + help="""When you specify multiple WHERE conditions on the command line +and let query-report build the query, join conditions with this operator. +Choices are {join_arg.choices_str()}. Default 'and'. +""") + cliutil.add_loglevel_argument(parser) + parser.add_argument( + 'query', + nargs=argparse.ZERO_OR_MORE, + help="""Query to run non-interactively. You can specify a full query +you write yourself, or conditions to follow WHERE and let query-report build +the rest of the query. +""") + args = parser.parse_args(arglist) + return args + +def main(arglist: Optional[Sequence[str]]=None, + stdout: TextIO=sys.stdout, + stderr: TextIO=sys.stderr, + config: Optional[configmod.Config]=None, +) -> int: + args = parse_arguments(arglist) + cliutil.set_loglevel(logger, args.loglevel) + if config is None: + config = configmod.Config() + config.load_file() + + fy = config.fiscal_year_begin() + if args.stop_date is None and args.start_date is not None: + args.stop_date = fy.next_fy_date(args.start_date) + query = build_query(args, fy, None if sys.stdin.isatty() else sys.stdin) + is_interactive = query is None and sys.stdin.isatty() + if args.report_type is None: + try: + args.report_type = ReportFormat[args.output_file.suffix[1:].upper()] + except (AttributeError, KeyError): + args.report_type = ReportFormat.TEXT # if is_interactive else ReportFormat.ODS + load_func = BooksLoader( + config.books_loader(), + args.start_date, + args.stop_date, + [rewrite.RewriteRuleset.from_yaml(path) for path in args.rewrite_rules], + ) + shell = BQLShell(is_interactive, load_func, stdout, args.report_type.value) + shell.on_Reload() + if query is None: + shell.cmdloop() + else: + shell.onecmd(query) + + return cliutil.ExitCode.OK + +entry_point = cliutil.make_entry_point(__name__, PROGNAME) + +if __name__ == '__main__': + exit(entry_point())