Fix UnboundLocalError when collect() is called before fit()

By deriving outcome and compevent model lists under a single guard with a None fallback

Author: remlapmot

pySEQTarget/SEQuential.py

@@ -453,13 +453,17 @@ def collect(self) -> SEQoutput:
             "collection_time": self._time_collected,
         }
 
-        if self.compevent_colname is not None:
-            compevent_models = [model["compevent"] for model in self.outcome_model]
-        else:
-            compevent_models = None
-
         if self.outcome_model is not None:
             outcome_models = [model["outcome"] for model in self.outcome_model]
+            if self.compevent_colname is not None:
+                compevent_models = [model["compevent"] for model in self.outcome_model]
+            else:
+                compevent_models = None
+        else:
+            # collect() before fit(): no models to report, but the rest of the
+            # output (diagnostics, timings) is still valid.
+            outcome_models = None
+            compevent_models = None
 
         if self.risk_estimates is None:
             risk_ratio = risk_difference = None

tests/test_accessor.py

@@ -31,6 +31,29 @@ def test_ITT_collector():
         collector.retrieve_data("unique_switches")
 
 
+def test_collect_before_fit():
+    # collect() without fit() must return an SEQoutput with None models rather
+    # than raising UnboundLocalError. Diagnostics from expand() still come
+    # through.
+    s = SEQuential(
+        load_data("SEQdata"),
+        id_col="ID",
+        time_col="time",
+        eligible_col="eligible",
+        treatment_col="tx_init",
+        outcome_col="outcome",
+        time_varying_cols=["N", "L", "P"],
+        fixed_cols=["sex"],
+        method="ITT",
+        parameters=SEQopts(),
+    )
+    s.expand()
+    collector = s.collect()
+    assert collector.outcome_models is None
+    assert collector.compevent_models is None
+    assert collector.retrieve_data("unique_outcomes").height > 0
+
+
 def test_censoring_collector_switch_diagnostics():
     # Under method="censoring" the switch diagnostics exist and must be
     # retrievable (regression for dict.has_key).