reports.core: Start RelatedPostings class.

conservancy_beancount/reports/core.py

@@ -0,0 +1,79 @@
+"""core.py - Common data classes for reporting functionality"""
+# Copyright © 2020  Brett Smith
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import collections
+
+from decimal import Decimal
+
+from .. import data
+
+from typing import (
+    overload,
+    Dict,
+    List,
+    Sequence,
+    Union,
+)
+
+class RelatedPostings(Sequence[data.Posting]):
+    """Collect and query related postings
+
+    This class provides common functionality for collecting related postings
+    and running queries on them: iterating over them, tallying their balance,
+    etc.
+
+    This class doesn't know anything about how the postings are related. That's
+    entirely up to the caller.
+
+    A common pattern is to use this class with collections.defaultdict
+    to organize postings based on some key::
+
+        report = collections.defaultdict(RelatedPostings)
+        for txn in transactions:
+            for post in Posting.from_txn(txn):
+                if should_report(post):
+                    key = post_key(post)
+                    report[key].add(post)
+    """
+
+    def __init__(self) -> None:
+        self._postings: List[data.Posting] = []
+
+    @overload
+    def __getitem__(self, index: int) -> data.Posting: ...
+
+    @overload
+    def __getitem__(self, s: slice) -> Sequence[data.Posting]: ...
+
+    def __getitem__(self,
+                    index: Union[int, slice],
+    ) -> Union[data.Posting, Sequence[data.Posting]]:
+        if isinstance(index, slice):
+            raise NotImplementedError("RelatedPostings[slice]")
+        else:
+            return self._postings[index]
+
+    def __len__(self) -> int:
+        return len(self._postings)
+
+    def add(self, post: data.Posting) -> None:
+        self._postings.append(post)
+
+    def balance(self) -> Sequence[data.Amount]:
+        currency_balance: Dict[str, Decimal] = collections.defaultdict(Decimal)
+        for post in self:
+            currency_balance[post.units.currency] += post.units.number
+        return [data.Amount(number, key) for key, number in currency_balance.items()]

tests/test_reports_related_postings.py

@@ -0,0 +1,87 @@
+"""test_reports_related_postings - Unit tests for RelatedPostings"""
+# Copyright © 2020  Brett Smith
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import datetime
+import itertools
+
+import pytest
+
+from . import testutil
+
+from conservancy_beancount import data
+from conservancy_beancount.reports import core
+
+_day_counter = itertools.count(1)
+def next_date():
+    return testutil.FY_MID_DATE + datetime.timedelta(next(_day_counter))
+
+def txn_pair(acct, src_acct, dst_acct, amount, date=None, txn_meta={}, post_meta={}):
+    if date is None:
+        date = next_date()
+    src_txn = testutil.Transaction(date=date, **txn_meta, postings=[
+        (acct, amount, post_meta.copy()),
+        (src_acct, -amount),
+    ])
+    dst_date = date + datetime.timedelta(days=1)
+    dst_txn = testutil.Transaction(date=dst_date, **txn_meta, postings=[
+        (acct, -amount, post_meta.copy()),
+        (dst_acct, amount),
+    ])
+    return (src_txn, dst_txn)
+
+def donation(amount, currency='USD', date=None, other_acct='Assets:Cash', **meta):
+    if date is None:
+        date = next_date()
+    return testutil.Transaction(date=date, postings=[
+        ('Income:Donations', -amount, currency, meta),
+        (other_acct, amount, currency),
+    ])
+
+def test_balance():
+    related = core.RelatedPostings()
+    related.add(data.Posting.from_beancount(donation(10), 0))
+    assert related.balance() == [testutil.Amount(-10)]
+    related.add(data.Posting.from_beancount(donation(15), 0))
+    assert related.balance() == [testutil.Amount(-25)]
+    related.add(data.Posting.from_beancount(donation(20), 0))
+    assert related.balance() == [testutil.Amount(-45)]
+
+def test_balance_zero():
+    related = core.RelatedPostings()
+    related.add(data.Posting.from_beancount(donation(10), 0))
+    related.add(data.Posting.from_beancount(donation(-10), 0))
+    assert related.balance() == [testutil.Amount(0)]
+
+def test_balance_multiple_currencies():
+    related = core.RelatedPostings()
+    related.add(data.Posting.from_beancount(donation(10, 'GBP'), 0))
+    related.add(data.Posting.from_beancount(donation(15, 'GBP'), 0))
+    related.add(data.Posting.from_beancount(donation(20, 'EUR'), 0))
+    related.add(data.Posting.from_beancount(donation(25, 'EUR'), 0))
+    assert set(related.balance()) == {
+        testutil.Amount(-25, 'GBP'),
+        testutil.Amount(-45, 'EUR'),
+    }
+
+def test_balance_multiple_currencies_one_zero():
+    related = core.RelatedPostings()
+    related.add(data.Posting.from_beancount(donation(10, 'EUR'), 0))
+    related.add(data.Posting.from_beancount(donation(15, 'USD'), 0))
+    related.add(data.Posting.from_beancount(donation(-10, 'EUR'), 0))
+    assert set(related.balance()) == {
+        testutil.Amount(-15, 'USD'),
+        testutil.Amount(0, 'EUR'),
+    }