diff --git a/conservancy_beancount/data.py b/conservancy_beancount/data.py index 351698a0769c47bd0e4a1ae471d2c5f8c0e2842b..73520bcb57985a9a1df55437c6d8782b7daa1bb6 100644 --- a/conservancy_beancount/data.py +++ b/conservancy_beancount/data.py @@ -40,6 +40,7 @@ from typing import ( ) from .beancount_types import ( + Directive, MetaKey, MetaValue, Posting as BasePosting, @@ -294,6 +295,19 @@ class Posting(BasePosting): for index, post in enumerate(txn.postings): yield cls.from_beancount(txn, index, post) + @classmethod + def from_entries(cls, entries: Iterable[Directive]) -> Iterable['Posting']: + """Yield an enhanced Posting object for every posting in these entries""" + for entry in entries: + # Because Beancount's own Transaction class isn't type-checkable, + # we can't statically check this. Might as well rely on duck + # typing while we're at it: just try to yield postings from + # everything, and ignore entries that lack a postings attribute. + try: + yield from cls.from_txn(entry) # type:ignore[arg-type] + except AttributeError: + pass + def balance_of(txn: Transaction, *preds: Callable[[Account], Optional[bool]], diff --git a/tests/test_data_posting.py b/tests/test_data_posting.py index b558530098c9017e0bea45db088da840b26ee938..9bd33473c38792d2d6b7474e2033d779bcae55c1 100644 --- a/tests/test_data_posting.py +++ b/tests/test_data_posting.py @@ -18,6 +18,8 @@ import pytest from . import testutil +from beancount.core import data as bc_data + from conservancy_beancount import data @pytest.fixture @@ -57,3 +59,26 @@ def test_from_txn(simple_txn): assert all(source[x] == post[x] for x in range(len(source) - 1)) assert isinstance(post.account, data.Account) assert post.meta['note'] # Only works with PostingMeta + +def test_from_entries_two_txns(simple_txn): + entries = [simple_txn, simple_txn] + sources = [post for txn in entries for post in txn.postings] + for source, post in zip(sources, data.Posting.from_entries(entries)): + assert all(source[x] == post[x] for x in range(len(source) - 1)) + assert isinstance(post.account, data.Account) + assert post.meta['note'] # Only works with PostingMeta + +def test_from_entries_mix_txns_and_other_directives(simple_txn): + meta = { + 'filename': __file__, + 'lineno': 75, + } + entries = [ + bc_data.Commodity(meta, testutil.FY_START_DATE, 'EUR'), + bc_data.Commodity(meta, testutil.FY_START_DATE, 'USD'), + simple_txn, + ] + for source, post in zip(simple_txn.postings, data.Posting.from_entries(entries)): + assert all(source[x] == post[x] for x in range(len(source) - 1)) + assert isinstance(post.account, data.Account) + assert post.meta['note'] # Only works with PostingMeta