diff --git a/conservancy_beancount/data.py b/conservancy_beancount/data.py index b985ef2..6fdab61 100644 --- a/conservancy_beancount/data.py +++ b/conservancy_beancount/data.py @@ -244,6 +244,18 @@ class Metadata(MutableMapping[MetaKey, MetaValue]): key, type(value).__name__, )) + @overload + def first_link(self, key: MetaKey, default: None=None) -> Optional[str]: ... + + @overload + def first_link(self, key: MetaKey, default: str) -> str: ... + + def first_link(self, key: MetaKey, default: Optional[str]=None) -> Optional[str]: + try: + return self.get_links(key)[0] + except (IndexError, TypeError): + return default + class PostingMeta(Metadata): """Combined access to posting metadata with its parent transaction metadata diff --git a/conservancy_beancount/reports/accrual.py b/conservancy_beancount/reports/accrual.py index ca3ea1f..879f0d8 100644 --- a/conservancy_beancount/reports/accrual.py +++ b/conservancy_beancount/reports/accrual.py @@ -220,11 +220,7 @@ class AccrualPostings(core.RelatedPostings): seen.add(entity) def first_links(self, key: MetaKey, default: Optional[str]=None) -> Iterator[Optional[str]]: - for post in self: - try: - yield post.meta.get_links(key)[0] - except (IndexError, TypeError): - yield default + return (post.meta.first_link(key, default) for post in self) def make_consistent(self) -> Iterator[Tuple[MetaValue, 'AccrualPostings']]: account_ok = isinstance(self.account, str) diff --git a/tests/test_data_metadata.py b/tests/test_data_metadata.py index 6fd0af9..effba9d 100644 --- a/tests/test_data_metadata.py +++ b/tests/test_data_metadata.py @@ -57,3 +57,35 @@ def test_get_links_bad_type(value): meta = data.Metadata({'key': value}) with pytest.raises(TypeError): meta.get_links('key') + +def test_first_link_from_txn(simple_txn): + meta = data.PostingMeta(simple_txn, 0) + assert meta.first_link('note') == 'txn' + +def test_first_link_from_post_override(simple_txn): + meta = data.PostingMeta(simple_txn, 1) + assert meta.first_link('note') == 'donation' + +def test_first_link_is_only_link(simple_txn): + meta = data.PostingMeta(simple_txn, 1) + assert meta.first_link('extra') == 'Extra' + +def test_first_link_nonexistent_metadata(simple_txn): + meta = data.PostingMeta(simple_txn, 1) + assert meta.first_link('Nonexistent') is None + +def test_first_link_nonexistent_default(simple_txn): + meta = data.PostingMeta(simple_txn, 1) + assert meta.first_link('Nonexistent', 'missing') == 'missing' + +@pytest.mark.parametrize('meta_value', testutil.NON_STRING_METADATA_VALUES) +def test_first_link_bad_type_metadata(simple_txn, meta_value): + simple_txn.meta['badmeta'] = meta_value + meta = data.PostingMeta(simple_txn, 1) + assert meta.first_link('badmeta') is None + +@pytest.mark.parametrize('meta_value', testutil.NON_STRING_METADATA_VALUES) +def test_first_link_bad_type_default(simple_txn, meta_value): + simple_txn.meta['badmeta'] = meta_value + meta = data.PostingMeta(simple_txn, 1) + assert meta.first_link('badmeta', '_') == '_'