From 396173b55d1bb62ac309d0d244c92fc960c2eaf7 Mon Sep 17 00:00:00 2001
From: Brett Smith <brettcsmith@brettcsmith.org>
Date: Thu, 21 May 2020 23:08:57 -0400
Subject: [PATCH] reports.Balance: Add eq_zero, ge_zero, and le_zero methods.

Support for RT#11294.
---
 conservancy_beancount/reports/core.py | 25 ++++++++++++--
 tests/test_reports_balance.py         | 50 +++++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 2 deletions(-)

diff --git a/conservancy_beancount/reports/core.py b/conservancy_beancount/reports/core.py
index a879299..cf38125 100644
--- a/conservancy_beancount/reports/core.py
+++ b/conservancy_beancount/reports/core.py
@@ -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):
diff --git a/tests/test_reports_balance.py b/tests/test_reports_balance.py
index bdba8ab..249a288 100644
--- a/tests/test_reports_balance.py
+++ b/tests/test_reports_balance.py
@@ -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},