diff options
Diffstat (limited to 'finance/flow.py')
| -rw-r--r-- | finance/flow.py | 31 |
1 files changed, 16 insertions, 15 deletions
diff --git a/finance/flow.py b/finance/flow.py index fb28228..a5817af 100644 --- a/finance/flow.py +++ b/finance/flow.py @@ -1,6 +1,7 @@ import dataclasses from datetime import datetime from decimal import Decimal +from typing import Generator @dataclasses.dataclass(kw_only=True, frozen=True) @@ -12,22 +13,22 @@ class Flow: def integrate(self, start: datetime, end: datetime) -> Decimal: """Integrate the flow between two dates to an amount of money""" - retval = Decimal(0.0) + payments: int = 0 - current = datetime(start.year, start.month, 1) - - if start == current: - retval += self.amount + for candidate in monthly_candidates(start): + if start <= candidate <= end: + payments += 1 + if candidate > end: + break - while True: - if current.month == 12: - current = datetime(current.year + 1, 1, 1) - else: - current = datetime(current.year, current.month + 1, 1) + return self.amount * Decimal(payments) - if current <= end: - retval += self.amount - else: - break - return retval +def monthly_candidates(start: datetime) -> Generator[datetime, None, None]: + current = datetime(start.year, start.month, 1) + while True: + yield current + if current.month == 12: + current = datetime(current.year + 1, 1, 1) + else: + current = datetime(current.year, current.month + 1, 1) |