conservancy_beancount/conservancy_beancount/reports/query.py

"""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())