diff --git a/conservancy_beancount/data.py b/conservancy_beancount/data.py index 29148d7d9bd9a59309dce29f9a94644a05be672f..458485c72af446eab587f410d712cf543e82051a 100644 --- a/conservancy_beancount/data.py +++ b/conservancy_beancount/data.py @@ -16,6 +16,8 @@ import collections.abc +from beancount.core import account as bc_account + from typing import ( Iterable, Iterator, @@ -29,6 +31,17 @@ from .beancount_types import ( Transaction, ) +class Account(str): + SEP = bc_account.sep + + def is_under(self, acct_s: str) -> bool: + return self.startswith(acct_s) and ( + acct_s.endswith(self.SEP) + or self == acct_s + or self[len(acct_s)] == self.SEP + ) + + class PostingMeta(collections.abc.MutableMapping): def __init__(self, txn: Transaction, index: int, post: Optional[Posting]=None) -> None: if post is None: diff --git a/tests/test_data_account.py b/tests/test_data_account.py new file mode 100644 index 0000000000000000000000000000000000000000..a3a619dcf139cc056b42c89501c67f41e1342fa3 --- /dev/null +++ b/tests/test_data_account.py @@ -0,0 +1,32 @@ +"""Test Account class""" +# 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 . + +import pytest + +from conservancy_beancount import data + +@pytest.mark.parametrize('acct_name,under_arg,expected', [ + ('Expenses:Tax:Sales', 'Expenses:Tax:Sales:', False), + ('Expenses:Tax:Sales', 'Expenses:Tax:Sales', True), + ('Expenses:Tax:Sales', 'Expenses:Tax:', True), + ('Expenses:Tax:Sales', 'Expenses:Tax', True), + ('Expenses:Tax:Sales', 'Expenses:', True), + ('Expenses:Tax:Sales', 'Expenses', True), + ('Expenses:Tax:Sales', 'Accrued:', False), + ('Expenses:Tax:Sales', 'Accrued', False), +]) +def test_is_under(acct_name, under_arg, expected): + assert data.Account(acct_name).is_under(under_arg) == expected