reports.Balance: Add eq_zero, ge_zero, and le_zero methods.

Support for RT#11294.
This commit is contained in:
Brett Smith 2020-05-21 23:08:57 -04:00
parent 1286cc2200
commit 396173b55d
2 changed files with 73 additions and 2 deletions

View file

@ -15,6 +15,7 @@
# along with this program. If not, see <https://www.gnu.org/licenses/>.
import collections
import operator
from decimal import Decimal
@ -22,6 +23,7 @@ from .. import data
from typing import (
overload,
Callable,
DefaultDict,
Dict,
Iterable,
@ -39,6 +41,8 @@ from ..beancount_types import (
MetaValue,
)
DecimalCompat = data.DecimalCompat
class Balance(Mapping[str, data.Amount]):
"""A collection of amounts mapped by currency
@ -81,8 +85,25 @@ class Balance(Mapping[str, data.Amount]):
def __len__(self) -> int:
return len(self._currency_map)
def is_zero(self) -> bool:
return all(number == 0 for number in self._currency_map.values())
def _all_amounts(self,
op_func: Callable[[DecimalCompat, DecimalCompat], bool],
operand: DecimalCompat,
) -> bool:
return all(op_func(number, operand) for number in self._currency_map.values())
def eq_zero(self) -> bool:
"""Returns true if all amounts in the balance == 0."""
return self._all_amounts(operator.eq, 0)
is_zero = eq_zero
def ge_zero(self) -> bool:
"""Returns true if all amounts in the balance >= 0."""
return self._all_amounts(operator.ge, 0)
def le_zero(self) -> bool:
"""Returns true if all amounts in the balance <= 0."""
return self._all_amounts(operator.le, 0)
class MutableBalance(Balance):

View file

@ -67,6 +67,56 @@ def test_mixed_balance():
assert not balance.is_zero()
assert all(balance[key] == amt for key, amt in amounts.items())
@pytest.mark.parametrize('balance_map_kwargs,expected', [
({}, True),
({'USD': 0}, True),
({'USD': 0, 'EUR': 0}, True),
({'USD': -10, 'EUR': 0}, False),
({'EUR': -10}, False),
({'USD': -10, 'EUR': -20}, False),
({'USD': 10, 'EUR': -20}, False),
({'JPY': 10}, False),
({'JPY': 10, 'BRL': 0}, False),
({'JPY': 10, 'BRL': 20}, False),
])
def test_eq_zero(balance_map_kwargs, expected):
amounts = testutil.balance_map(**balance_map_kwargs)
balance = core.Balance(amounts.items())
assert balance.eq_zero() == expected
assert balance.is_zero() == expected
@pytest.mark.parametrize('balance_map_kwargs,expected', [
({}, True),
({'USD': 0}, True),
({'USD': 0, 'EUR': 0}, True),
({'EUR': -10}, False),
({'USD': 10, 'EUR': -20}, False),
({'USD': -10, 'EUR': -20}, False),
({'JPY': 10}, True),
({'JPY': 10, 'BRL': 0}, True),
({'JPY': 10, 'BRL': 20}, True),
])
def test_ge_zero(balance_map_kwargs, expected):
amounts = testutil.balance_map(**balance_map_kwargs)
balance = core.Balance(amounts.items())
assert balance.ge_zero() == expected
@pytest.mark.parametrize('balance_map_kwargs,expected', [
({}, True),
({'USD': 0}, True),
({'USD': 0, 'EUR': 0}, True),
({'EUR': -10}, True),
({'USD': 10, 'EUR': -20}, False),
({'USD': -10, 'EUR': -20}, True),
({'JPY': 10}, False),
({'JPY': 10, 'BRL': 0}, False),
({'JPY': 10, 'BRL': 20}, False),
])
def test_le_zero(balance_map_kwargs, expected):
amounts = testutil.balance_map(**balance_map_kwargs)
balance = core.Balance(amounts.items())
assert balance.le_zero() == expected
@pytest.mark.parametrize('balance_map_kwargs', [
{},
{'USD': 0},