diff --git a/fin_defs/__init__.py b/fin_defs/__init__.py index 6f8286e..dc716c4 100644 --- a/fin_defs/__init__.py +++ b/fin_defs/__init__.py @@ -505,6 +505,11 @@ class AssetInformation: THREE_DECIMALS_UNDER_THIS_AMOUNT = 0.10 +def assert_same_asset(do: str, a_asset: 'AssetAmount', b_asset: 'AssetAmount'): + if a_asset.asset != b_asset.asset: + msg = f'Attempting to {do} {a_asset} and {b_asset}, but these must have the same asset' + raise ValueError(msg) + @enforce_typing.enforce_types @dataclasses.dataclass(frozen=True, eq=True, slots=True) @@ -547,11 +552,7 @@ class AssetAmount: return other if other.is_zero(): return self - if self.asset != other.asset: - msg = ( - f'AssetAmount must have same asset, but: {self.asset} != {other.asset}' - ) - raise ValueError(msg) + assert_same_asset('add', self, other) return AssetAmount(self.asset, self.amount + other.amount) @@ -564,18 +565,13 @@ class AssetAmount: return -other if other.is_zero(): return self - if self.asset != other.asset: - msg = ( - f'AssetAmount must have same asset, but: {self.asset} != {other.asset}' - ) - raise ValueError(msg) + assert_same_asset('subtract', self, other) return AssetAmount(self.asset, self.amount - other.amount) + def __truediv__(self, other: 'Decimal | AssetAmount') -> 'Decimal | AssetAmount': if isinstance(other, AssetAmount): - if self.asset != other.asset: - msg = f'Attempting to divide {self} by {other}, but these must have the same asset' - raise ValueError(msg) + assert_same_asset('divide', self, other) return self.amount / other.amount return AssetAmount(self.asset, self.amount / other) @@ -588,11 +584,7 @@ class AssetAmount: raise TypeError(msg) if self.is_zero() or other.is_zero(): return self.amount - other.amount - if self.asset != other.asset: - msg = ( - f'AssetAmount must have same asset, but: {self.asset} != {other.asset}' - ) - raise ValueError(msg) + assert_same_asset('compare', self, other) return self.amount - other.amount def is_zero(self) -> bool: diff --git a/test/test_asset_amount.py b/test/test_asset_amount.py index 5c24961..6c988a6 100644 --- a/test/test_asset_amount.py +++ b/test/test_asset_amount.py @@ -1,4 +1,5 @@ from decimal import Decimal +import re import pytest @@ -83,7 +84,7 @@ def test_add_wrong_type(): def test_add_wrong_asset(): with pytest.raises( ValueError, - match='AssetAmount must have same asset, but: fiat:USD != fiat:DKK', + match=re.escape('Attempting to add $20.00 USD and 21.00 DKK, but these must have the same asset'), ): assert USD_20 + DKK_21 @@ -96,7 +97,7 @@ def test_sub_wrong_type(): def test_sub_wrong_asset(): with pytest.raises( ValueError, - match='AssetAmount must have same asset, but: fiat:USD != fiat:DKK', + match=re.escape('Attempting to subtract $20.00 USD and 21.00 DKK, but these must have the same asset'), ): assert USD_20 - DKK_21 @@ -109,7 +110,7 @@ def test_cmp_wrong_type(): def test_cmp_wrong_asset(): with pytest.raises( ValueError, - match='AssetAmount must have same asset, but: fiat:USD != fiat:DKK', + match=re.escape('Attempting to compare $20.00 USD and 21.00 DKK, but these must have the same asset'), ): assert USD_20 < DKK_21 @@ -117,6 +118,6 @@ def test_cmp_wrong_asset(): def test_div_wrong_asset(): with pytest.raises( ValueError, - match='AssetAmount must have same asset, but: fiat:USD != fiat:DKK', + match=re.escape('Attempting to divide $21.00 USD and 21.00 DKK, but these must have the same asset'), ): assert USD_21 / DKK_21