refactor: update rounded_survival_table to use dynamic column names

acro/acro_tables.py

@@ -522,7 +522,7 @@ def survival_plot(  # pylint: disable=too-many-arguments
     ):
         """Create the survival plot according to the status of suppressing."""
         if self.suppress:
-            survival_table = rounded_survival_table(survival_table)
+            survival_table = _rounded_survival_table(survival_table)
             plot = survival_table.plot(y="rounded_survival_fun", xlim=0, ylim=0)
         else:  # pragma: no cover
             plot = survival_func.plot()
@@ -914,14 +914,37 @@ def delete_empty_rows_columns(table: DataFrame) -> tuple[DataFrame, list[str]]:
     return (table, comments)
 
 
-def rounded_survival_table(survival_table):
-    """Calculate the rounded surival function."""
+def _rounded_survival_table(
+    survival_table: pd.DataFrame,
+    num_at_risk_col: str = "num at risk",
+    num_events_col: str = "num events",
+) -> pd.DataFrame:
+    """Calculate the rounded survival function.
+
+    Internal helper function for survival analysis with disclosure control.
+    Applies rounding to survival tables to prevent disclosure of small counts.
+
+    Parameters
+    ----------
+    survival_table : pd.DataFrame
+        The survival table containing survival analysis results.
+    num_at_risk_col : str, default "num at risk"
+        Name of the column containing number at risk values.
+    num_events_col : str, default "num events"
+        Name of the column containing number of events.
+
+    Returns
+    -------
+    pd.DataFrame
+        The survival table with rounded survival function added.
+    """
     death_censored = (
-        survival_table["num at risk"].shift(periods=1) - survival_table["num at risk"]
+        survival_table[num_at_risk_col].shift(periods=1)
+        - survival_table[num_at_risk_col]
     )
     death_censored = death_censored.tolist()
-    survivor = survival_table["num at risk"].tolist()
-    deaths = survival_table["num events"].tolist()
+    survivor = survival_table[num_at_risk_col].tolist()
+    deaths = survival_table[num_events_col].tolist()
     rounded_num_of_deaths = []
     rounded_num_at_risk = []
     sub_total = 0

test/test_initial.py

@@ -11,7 +11,7 @@
 import statsmodels.api as sm
 
 from acro import ACRO, acro_tables, add_constant, add_to_acro, record, utils
-from acro.acro_tables import rounded_survival_table
+from acro.acro_tables import _rounded_survival_table
 from acro.record import Records, load_records
 
 # pylint: disable=redefined-outer-name,too-many-lines
@@ -735,7 +735,7 @@ def test_rounded_survival_table():
     )
 
     # Apply rounded_survival_table
-    result = rounded_survival_table(survival_table.copy())
+    result = _rounded_survival_table(survival_table.copy())
 
     # Check that it has the rounded_survival_fun column
     assert "rounded_survival_fun" in result.columns