diff --git a/conservancy_beancount/reports/query.py b/conservancy_beancount/reports/query.py index 77a522212c1533732b0693033a9afc229e9d77ef..9ca96113e0960d15187257d67790bb4d0c64135b 100644 --- a/conservancy_beancount/reports/query.py +++ b/conservancy_beancount/reports/query.py @@ -33,6 +33,7 @@ from typing import ( Union, ) from ..beancount_types import ( + MetaKey, MetaValue, Posting, Transaction, @@ -41,6 +42,8 @@ from ..beancount_types import ( from decimal import Decimal from pathlib import Path from beancount.core.amount import _Amount as BeancountAmount +from beancount.core.inventory import Inventory +from beancount.core.position import _Position as Position import beancount.query.numberify as bc_query_numberify import beancount.query.query_compile as bc_query_compile @@ -59,10 +62,6 @@ from .. import config as configmod from .. import data from .. import rtutil -BUILTIN_FIELDS: AbstractSet[str] = frozenset(itertools.chain( - bc_query_env.TargetsEnvironment.columns, # type:ignore[has-type] - bc_query_env.TargetsEnvironment.functions, # type:ignore[has-type] -)) PROGNAME = 'query-report' logger = logging.getLogger('conservancy_beancount.reports.query') @@ -80,6 +79,12 @@ EnvironmentFunctions = Dict[ RowTypes = Sequence[Tuple[str, Type]] Rows = Sequence[NamedTuple] Store = List[Any] +QueryExpression = Union[ + bc_query_parser.Column, + bc_query_parser.Constant, + bc_query_parser.Function, + bc_query_parser.UnaryOp, +] QueryStatement = Union[ bc_query_parser.Balances, bc_query_parser.Journal, @@ -124,14 +129,45 @@ class BooksLoader: class QueryODS(core.BaseODS[NamedTuple, None]): + META_FNAMES = frozenset([ + 'any_meta', + 'entry_meta', + 'meta', + 'meta_docs', + 'str_meta', + ]) + def is_empty(self) -> bool: return not self.sheet.childNodes def section_key(self, row: NamedTuple) -> None: return None + def _generic_cell(self, value: Any) -> odf.table.TableCell: + if isinstance(value, Iterable) and not isinstance(value, (str, tuple)): + return self.multiline_cell(value) + else: + return self.string_cell('' if value is None else str(value)) + + def _inventory_cell(self, value: Inventory) -> odf.table.TableCell: + return self.balance_cell(core.Balance(pos.units for pos in value)) + + def _link_string_cell(self, value: str) -> odf.table.TableCell: + return self.meta_links_cell(value.split()) + + def _metadata_cell(self, value: MetaValue) -> odf.table.TableCell: + return self._cell_type(type(value))(value) + + def _position_cell(self, value: Position) -> odf.table.TableCell: + return self.currency_cell(value.units) + def _cell_type(self, row_type: Type) -> CellFunc: - if issubclass(row_type, BeancountAmount): + """Return a function to create a cell, for non-metadata row types.""" + if issubclass(row_type, Inventory): + return self._inventory_cell + elif issubclass(row_type, Position): + return self._position_cell + elif issubclass(row_type, BeancountAmount): return self.currency_cell elif issubclass(row_type, (int, float, Decimal)): return self.float_cell @@ -142,49 +178,85 @@ class QueryODS(core.BaseODS[NamedTuple, None]): else: return self._generic_cell - def _generic_cell(self, value: Any) -> odf.table.TableCell: - return self.string_cell('' if value is None else str(value)) - - def _link_cell(self, value: MetaValue) -> odf.table.TableCell: - if isinstance(value, str): - return self.meta_links_cell(value.split()) + def _link_cell_type(self, row_type: Type) -> CellFunc: + """Return a function to create a cell from metadata with documentation links.""" + if issubclass(row_type, str): + return self._link_string_cell + elif issubclass(row_type, tuple): + return self._generic_cell + elif issubclass(row_type, Iterable): + return self.meta_links_cell else: - return self._generic_cell(value) - - def _metadata_cell(self, value: MetaValue) -> odf.table.TableCell: - return self._cell_type(type(value))(value) + return self._generic_cell - def _cell_types(self, row_types: RowTypes) -> Iterator[CellFunc]: - for name, row_type in row_types: - if row_type is object: - if name.replace('_', '-') in data.LINK_METADATA: - yield self._link_cell - else: - yield self._metadata_cell - else: + def _meta_target(self, target: QueryExpression) -> Optional[MetaKey]: + """Return the metadata key looked up by this target, if any + + This function takes a parsed target (i.e., what we're SELECTing) and + recurses it to see whether it's looking up any metadata. If so, it + returns the key of that metadata. Otherwise it returns None. + """ + if isinstance(target, bc_query_parser.UnaryOp): + return self._meta_target(target.operand) + elif not isinstance(target, bc_query_parser.Function): + return None + try: + operand = target.operands[0] + except IndexError: + return None + if (target.fname in self.META_FNAMES + and isinstance(operand, bc_query_parser.Constant)): + return operand.value # type:ignore[no-any-return] + else: + for operand in target.operands: + retval = self._meta_target(operand) + if retval is not None: + break + return retval + + def _cell_types(self, statement: QueryStatement, row_types: RowTypes) -> Iterator[CellFunc]: + """Return functions to create table cells from result rows + + Given a parsed query and the types of return rows, yields a function + to create a cell for each column in the row, in order. The returned + functions vary in order to provide the best available formatting for + different data types. + """ + if (isinstance(statement, bc_query_parser.Select) + and isinstance(statement.targets, Sequence)): + targets = [t.expression for t in statement.targets] + else: + # Synthesize something that makes clear we're not loading metadata. + targets = [bc_query_parser.Column(name) for name, _ in row_types] + for target, (_, row_type) in zip(targets, row_types): + meta_key = self._meta_target(target) + if meta_key is None: yield self._cell_type(row_type) + elif meta_key in data.LINK_METADATA: + yield self._link_cell_type(row_type) + else: + yield self._metadata_cell - def write_query(self, row_types: RowTypes, rows: Rows) -> None: + def write_query(self, statement: QueryStatement, row_types: RowTypes, rows: Rows) -> None: if self.is_empty(): self.sheet.setAttribute('name', "Query 1") else: self.use_sheet(f"Query {len(self.document.spreadsheet.childNodes) + 1}") for name, row_type in row_types: - if row_type is object or issubclass(row_type, str): - col_width = 2.0 - elif issubclass(row_type, BeancountAmount): + if issubclass(row_type, datetime.date): + col_width = 1.0 + elif issubclass(row_type, (BeancountAmount, Inventory, Position)): col_width = 1.5 else: - col_width = 1.0 + col_width = 2.0 col_style = self.column_style(col_width) self.sheet.addElement(odf.table.TableColumn(stylename=col_style)) self.add_row(*( - self.string_cell(data.Metadata.human_name(name.replace('_', '-')), - stylename=self.style_bold) + self.string_cell(data.Metadata.human_name(name), stylename=self.style_bold) for name, _ in row_types )) self.lock_first_row() - cell_funcs = list(self._cell_types(row_types)) + cell_funcs = list(self._cell_types(statement, row_types)) for row in rows: self.add_row(*( cell_func(value) @@ -238,7 +310,7 @@ class AggregateSet(bc_query_compile.EvalAggregator): def update(self, store: Store, context: Context) -> None: value, = self.eval_args(context) - if isinstance(value, Sequence) and not isinstance(value, str): + if isinstance(value, Sequence) and not isinstance(value, (str, tuple)): store[self.handle].update(value) else: store[self.handle].add(value) @@ -304,9 +376,9 @@ class BQLShell(bc_query_shell.BQLShell): print("(empty)", file=self.outfile) else: logger.debug("rendering query as %s", output_format) - render_func(row_types, rows) + render_func(statement, row_types, rows) - def _render_csv(self, row_types: RowTypes, rows: Rows) -> None: + def _render_csv(self, statement: QueryStatement, row_types: RowTypes, rows: Rows) -> None: bc_query_render.render_csv( row_types, rows, @@ -315,11 +387,15 @@ class BQLShell(bc_query_shell.BQLShell): self.vars['expand'], ) - def _render_ods(self, row_types: RowTypes, rows: Rows) -> None: - self.ods.write_query(row_types, rows) - logger.info("results saved in sheet %s", self.ods.sheet.getAttribute('name')) + def _render_ods(self, statement: QueryStatement, row_types: RowTypes, rows: Rows) -> None: + self.ods.write_query(statement, row_types, rows) + logger.info( + "%s rows of results saved in sheet %s", + len(rows), + self.ods.sheet.getAttribute('name'), + ) - def _render_text(self, row_types: RowTypes, rows: Rows) -> None: + def _render_text(self, statement: QueryStatement, row_types: RowTypes, rows: Rows) -> None: with contextlib.ExitStack() as stack: if self.is_interactive: output = stack.enter_context(self.get_pager()) @@ -394,9 +470,7 @@ ODS reports. help="""Query to run non-interactively. If none is provided, and standard input is not a terminal, reads the query from stdin instead. """) - - args = parser.parse_args(arglist) - return args + return parser.parse_args(arglist) def main(arglist: Optional[Sequence[str]]=None, stdout: TextIO=sys.stdout,