Curo code refactoring

Author: andrewmurphy353

curo/calculator.py

@@ -2,7 +2,6 @@
 Docstring for calculator
 """
 
-from dataclasses import dataclass
 from typing import List, Optional
 import pandas as pd
 from scipy.optimize import root_scalar
@@ -13,7 +12,6 @@
 from curo.exceptions import UnsolvableError, ValidationError
 from curo.utils import roll_date, gauss_round, to_timestamp
 
-@dataclass
 class Calculator:
     """
     The Calculator class provides the entry point for solving unknown values and/or
@@ -29,7 +27,7 @@ class Calculator:
     Attributes:
         precision (int): The rounding precision for cash flow values.
         profile (pd.DataFrame or None): The DataFrame containing cash flow data.
-        series (List[Series]): The list of cash flow series (empty if not provided).
+        _series (List[Series]): Private list of provided cash flow series.
         _is_bespoke_profile (bool): Used internally to identify profile source.
 
     Raises:
@@ -39,15 +37,13 @@ class Calculator:
         Calculator: A new Calculator instance configured with the specified precision
             and optional profile.
     """
-    precision: int = 2
-    profile: Optional[pd.DataFrame] = None
-    series: List[Series] = None
-
-    def __post_init__(self):
-        if not 0 <= self.precision <= 4:
+    def __init__(self, precision: int = 2, profile: Optional[pd.DataFrame] = None):
+        if not 0 <= precision <= 4:
             raise ValidationError("Precision must be between 0 and 4")
-        self.series = self.series or []
-        self._is_bespoke_profile = self.profile is not None
+        self.precision = precision
+        self.profile = profile
+        self._series = []
+        self._is_bespoke_profile = profile is not None
 
     def add(self, series: Series) -> None:
         """
@@ -75,7 +71,7 @@ def add(self, series: Series) -> None:
         if series.amount is not None:
             # Coerce series monetary value to specified precision
             series.amount = gauss_round(series.amount, self.precision)
-        self.series.append(series)
+        self._series.append(series)
 
     def solve_value(
         self,
@@ -242,7 +238,7 @@ def _build_profile(self, start_date: Optional[pd.Timestamp] = None) -> pd.DataFr
         advance_start_date = start_date
         payment_start_date = start_date
         charge_start_date = start_date
-        for s in self.series:
+        for s in self._series:
             if isinstance(s, SeriesAdvance):
                 if s.post_date_from is None and s.mode == Mode.ARREAR:
                     advance_start_date = roll_date(
@@ -314,7 +310,7 @@ def _build_profile(self, start_date: Optional[pd.Timestamp] = None) -> pd.DataFr
 
         if not cash_flows_list:
             # Handle empty case
-            return pd.DataFrame(columns=[
+            profile = pd.DataFrame(columns=[
                 Column.POST_DATE.value,
                 Column.VALUE_DATE.value,
                 Column.AMOUNT.value,
@@ -325,6 +321,17 @@ def _build_profile(self, start_date: Optional[pd.Timestamp] = None) -> pd.DataFr
                 Column.IS_INTEREST_CAPITALISED.value,
                 Column.IS_CHARGE.value
             ])
+            return profile.astype({
+                Column.POST_DATE.value: 'datetime64[ns, UTC]',
+                Column.VALUE_DATE.value: 'datetime64[ns, UTC]',
+                Column.AMOUNT.value: 'float64',
+                Column.IS_KNOWN.value: 'bool',
+                Column.WEIGHTING.value: 'float64',
+                Column.LABEL.value: 'object',
+                Column.MODE.value: 'object',
+                Column.IS_INTEREST_CAPITALISED.value: 'object',
+                Column.IS_CHARGE.value: 'bool'
+            })
 
         profile = pd.concat(cash_flows_list, ignore_index=True)
         # Ensure consistent dtypes
@@ -564,11 +571,6 @@ def _validate_profile(
                 raise ValidationError(
                     "Unknowns must be either all advances or all payments, not both"
                 )
-            # Check weighting > 0 for unknowns
-            if unknowns[Column.WEIGHTING.value].le(0).any():
-                raise ValidationError(
-                    "Unknown values must have positive weighting in SOLVE_VALUE mode"
-                )
         elif mode == ValidationMode.SOLVE_RATE:
             if not unknowns.empty:
                 raise ValidationError("All values must be known in SOLVE_RATE mode")
@@ -585,32 +587,64 @@ def _assign_factors(self, cash_flows: pd.DataFrame, day_count: Convention) -> pd
 
         Args:
             cash_flows (pd.DataFrame): DataFrame containing cash flow data with
-                CashFlowColumn.POST_DATE.
+                CashFlowColumn.POST_DATE and CashFlowColumn.VALUE_DATE.
             day_count (Convention): The day count convention for computing time intervals.
 
         Returns:
             pd.DataFrame: A copy of the input DataFrame with an added 'factor' column
-            containing computed time factors.
+                containing computed DayCountFactor objects.
 
         Notes:
-        - Factors are computed from the first `post_date` (if DayCountOrigin.DRAWDOWN)
-          or between consecutive `post_date` values (otherwise).
-        - The 'factor' column contains objects specific to the day count convention.
+            - Factors are computed from the first advance's post_date or value_date (if
+              DayCountOrigin.DRAWDOWN) or between consecutive dates (if DayCountOrigin.NEIGHBOUR).
+            - For charges, if include_non_financing_flows is False, the factor is computed
+              between the same date (zero period).
+            - The date used (post_date or value_date) depends on day_count.use_post_dates.
         """
         cash_flows = cash_flows.copy()
-        if day_count.day_count_origin == DayCountOrigin.DRAWDOWN:
-            start_date = cash_flows['post_date'].iloc[0]
-            cash_flows['factor'] = cash_flows['post_date'].apply(
-                lambda d: day_count.compute_factor(start_date, d)
-            )
-        else:
-            cash_flows['prev_date'] = cash_flows['post_date'].shift(1)
-            cash_flows.loc[0, 'prev_date'] = cash_flows.loc[0, 'post_date']
-            cash_flows['factor'] = cash_flows.apply(
-                lambda row: day_count.compute_factor(row['prev_date'], row['post_date']),
-                axis=1
-            )
-            cash_flows = cash_flows.drop(columns=['prev_date'])
+        date_column = (
+            Column.POST_DATE.value
+            if day_count.use_post_dates
+            else Column.VALUE_DATE.value
+        )
+
+        # Initialize the factor column as object type to store DayCountFactor
+        cash_flows['factor'] = None
+
+        # Find the first advance's date for DRAWDOWN origin
+        advances = cash_flows[
+            (~cash_flows[Column.IS_CHARGE.value]) &
+            (cash_flows[Column.IS_INTEREST_CAPITALISED.value].isna())
+        ]
+        if advances.empty:
+            raise ValidationError("At least one advance required for factor assignment")
+        drawdown_date = advances[date_column].min()
+
+        # For NEIGHBOUR origin, track the previous date
+        neighbour_date = drawdown_date
+
+        for idx in cash_flows.index:
+            cash_flow_date = cash_flows.loc[idx, date_column]
+            is_charge = cash_flows.loc[idx, Column.IS_CHARGE.value]
+
+            # Handle charges when non-financing flows are excluded
+            if is_charge and not day_count.include_non_financing_flows:
+                factor = day_count.compute_factor(cash_flow_date, cash_flow_date)
+                cash_flows.at[idx, 'factor'] = factor
+                continue
+
+            # Handle cash flows predating or equal to drawdown_date
+            if cash_flow_date <= drawdown_date:
+                factor = day_count.compute_factor(cash_flow_date, cash_flow_date)
+            else:
+                if day_count.day_count_origin == DayCountOrigin.DRAWDOWN:
+                    factor = day_count.compute_factor(drawdown_date, cash_flow_date)
+                else:  # NEIGHBOUR
+                    factor = day_count.compute_factor(neighbour_date, cash_flow_date)
+                    neighbour_date = cash_flow_date
+
+            cash_flows.at[idx, 'factor'] = factor
+
         return cash_flows
 
     def _calculate_nfv(self, cash_flows: pd.DataFrame, day_count: Convention, rate: float) -> float:

tests/test_calculator.py

@@ -1 +1,86 @@
-# general tests - constructor, add method
+# test_calculator.py
+# pylint: disable=C0114,C0116,W0212,W0621
+# - C0114: missing-module-docstring
+# - C0116: missing-function-docstring
+# - W0212: protected-access
+# - W0621: redefined-outer-name
+
+import pytest
+import pandas as pd
+from curo.calculator import Calculator
+from curo.series import SeriesAdvance, SeriesPayment, SeriesCharge
+from curo.enums import CashFlowColumn as Column
+from curo.exceptions import ValidationError
+
+@pytest.fixture
+def valid_series():
+    return SeriesAdvance(
+        number_of=1,
+        amount=1000.0,
+        post_date_from=pd.Timestamp("2025-01-01", tz="UTC"),
+        label="Loan advance"
+    )
+
+@pytest.fixture
+def valid_profile():
+    return pd.DataFrame({
+        Column.POST_DATE.value: [pd.Timestamp("2025-01-01", tz="UTC")],
+        Column.VALUE_DATE.value: [pd.Timestamp("2025-01-01", tz="UTC")],
+        Column.AMOUNT.value: [1000.0],
+        Column.IS_KNOWN.value: [True],
+        Column.WEIGHTING.value: [1.0],
+        Column.LABEL.value: ["Loan advance"],
+        Column.MODE.value: ["advance"],
+        Column.IS_INTEREST_CAPITALISED.value: [None],
+        Column.IS_CHARGE.value: [False]
+    })
+
+def test_constructor_valid_precision():
+    calc = Calculator(precision=2)
+    assert calc.precision == 2
+    assert calc.profile is None
+    assert not calc._series
+    assert calc._is_bespoke_profile is False
+
+def test_constructor_invalid_precision():
+    with pytest.raises(ValidationError, match="Precision must be between 0 and 4"):
+        Calculator(precision=5)
+    with pytest.raises(ValidationError, match="Precision must be between 0 and 4"):
+        Calculator(precision=-1)
+
+def test_constructor_with_profile(valid_profile):
+    calc = Calculator(precision=3, profile=valid_profile)
+    assert calc.precision == 3
+    assert calc.profile.equals(valid_profile)
+    assert not calc._series
+    assert calc._is_bespoke_profile is True
+
+def test_add_valid_series(valid_series):
+    calc = Calculator(precision=2)
+    calc.add(valid_series)
+    assert len(calc._series) == 1
+    assert calc._series[0] == valid_series
+    assert calc._series[0].amount == 1000.0  # Amount unchanged
+
+def test_add_series_with_rounding():
+    calc = Calculator(precision=1)
+    series = SeriesPayment(
+        number_of=1,
+        amount=123.456,
+        post_date_from=pd.Timestamp("2025-01-01", tz="UTC"),
+        label="Loan repayment"
+    )
+    calc.add(series)
+    assert len(calc._series) == 1
+    assert calc._series[0].amount == 123.5  # Rounded to 1 decimal place
+
+def test_add_with_bespoke_profile(valid_profile):
+    calc = Calculator(precision=2, profile=valid_profile)
+    series = SeriesCharge(
+        number_of=1,
+        amount=50.0,
+        post_date_from=pd.Timestamp("2025-01-01", tz="UTC"),
+        label="Arrangement fee"
+    )
+    with pytest.raises(ValidationError, match="Cannot add series with a bespoke profile"):
+        calc.add(series)