''' Utilities for custom importers. ''' __copyright__ = 'Copyright (c) 2020 Kristóf Marussy ' __license__ = 'GNU GPLv2' from abc import ABC, abstractmethod import datetime as dt from decimal import Decimal from typing import cast, Callable, Iterable, List, NamedTuple, Optional, Set, TypeVar, Union from beancount.core import amount as am, data from beancount.core.amount import Amount from beancount.core.flags import FLAG_OKAY, FLAG_WARNING from beancount.core.number import D, ZERO MISSING_AMOUNT = cast(Amount, None) COMMENT_META = 'import-raw-comment' PAYEE_META = 'import-raw-payee' class InvalidEntry(Exception): pass class Posting(NamedTuple): account: str amount: Amount class Row(ABC): entry_type: Optional[str] payee: Optional[str] comment: str meta: data.Meta flag: str tags: Set[str] links: Set[str] _postings: Optional[List[Posting]] def __init__(self, file_name: str, line_number: int, entry_type: Optional[str], payee: Optional[str], comment: str): self.entry_type = entry_type self.payee = payee self.comment = comment self.meta = data.new_metadata(file_name, line_number) self.flag = FLAG_OKAY self.tags = set() self.links = set() self._postings = None @property @abstractmethod def transacted_amount(self) -> Amount: pass @property def transacted_currency(self) -> str: return self.transacted_amount.currency @property def postings(self) -> Optional[List[Posting]]: return self._postings def assign_to_accounts(self, *postings: Posting) -> None: if self.done: raise InvalidEntry('Transaction is alrady done processing') self._postings = list(postings) if not self._postings: raise InvalidEntry('Not assigned to any accounts') head, *rest = self._postings sum = head.amount for posting in rest: sum = am.add(sum, posting.amount) if sum != self.transacted_amount: self.flag = FLAG_WARNING def assign_to_account(self, account: str) -> None: self.assign_to_accounts(Posting(account, self.transacted_amount)) @property def done(self) -> bool: return self._postings is not None Extractor = Callable[[Row], None] TRow = TypeVar('TRow', bound=Row) def run_row_extractors(row: TRow, extractors: Iterable[Callable[[TRow], None]]) -> None: for extractor in extractors: extractor(row) if row.done: return def extract_unknown(expenses_account: str, income_account: str) -> Extractor: def do_extract(row: Row) -> None: if row.transacted_amount.number < ZERO: row.assign_to_account(expenses_account) else: row.assign_to_account(income_account) row.flag = FLAG_WARNING return do_extract def parse_date(date_str: str, format_string: str) -> dt.date: try: return dt.datetime.strptime(date_str, format_string).date() except ValueError as exc: raise InvalidEntry(f'Cannot parse date: {date_str}') from exc def parse_number(in_amount: Union[str, int, float, Decimal]) -> Decimal: try: value = D(in_amount) except ValueError as exc: raise InvalidEntry(f'Cannot parse number: {in_amount}') from exc if value is None: raise InvalidEntry(f'Parse number returned None: {in_amount}') return value