import dataclasses from datetime import datetime from decimal import Decimal from typing import Generator @dataclasses.dataclass(kw_only=True, frozen=True) class Flow: """Time-discrete flow of money paid on the first day of a month""" amount: Decimal since: None | datetime until: None | datetime def integrate(self, start: datetime, end: datetime) -> Decimal: """Integrate the flow between two dates to an amount of money""" payments: int = 0 if self.since is not None: if start < self.since: start = self.since if self.until is not None: if end > self.until: end = self.until for candidate in monthly_candidates(start): if start <= candidate <= end: payments += 1 if candidate > end: break return self.amount * Decimal(payments) 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)