generalize this

Author: CamDavidsonPilon

core/pioreactor/estimators/__init__.py

@@ -21,7 +21,7 @@ def _estimator_path_for(device: str, name: str) -> Path:
     return ESTIMATOR_PATH / device / f"{name}.yaml"
 
 
-def load_active_estimator(device: Device) -> structs.ODFusionEstimator | None:
+def load_active_estimator(device: Device) -> structs.AnyEstimator | None:
     with local_persistent_storage("active_estimators") as storage:
         active_name = storage.get(device)
 
@@ -30,15 +30,15 @@ def load_active_estimator(device: Device) -> structs.ODFusionEstimator | None:
     return load_estimator(device, active_name)
 
 
-def load_estimator(device: Device, estimator_name: str) -> structs.ODFusionEstimator:
+def load_estimator(device: Device, estimator_name: str) -> structs.AnyEstimator:
     target_file = _estimator_path_for(device, estimator_name)
     if not target_file.is_file():
         raise FileNotFoundError(f"Estimator {estimator_name} was not found in {ESTIMATOR_PATH / device}")
     if target_file.stat().st_size == 0:
         raise FileNotFoundError(f"Estimator {estimator_name} is empty")
 
     try:
-        return yaml_decode(target_file.read_bytes(), type=structs.ODFusionEstimator)
+        return yaml_decode(target_file.read_bytes(), type=structs.subclass_union(structs.EstimatorBase))
     except ValidationError as exc:
         raise ValidationError(f"Error reading {target_file.stem}: {exc}") from exc
 

core/pioreactor/web/tasks.py

@@ -551,14 +551,15 @@ def calibration_save_calibration(device: str, calibration_payload: dict[str, obj
 def estimator_save_estimator(device: str, estimator_payload: dict[str, object]) -> dict[str, str]:
     from msgspec.json import decode as json_decode
     from msgspec.json import encode as json_encode
-    from pioreactor.structs import ODFusionEstimator
+    from pioreactor.structs import EstimatorBase
+    from pioreactor.structs import subclass_union
 
     logger.debug(
         "Starting estimator save: device=%s payload_keys=%s",
         device,
         sorted(estimator_payload.keys()),
     )
-    estimator = json_decode(json_encode(estimator_payload), type=ODFusionEstimator)
+    estimator = json_decode(json_encode(estimator_payload), type=subclass_union(EstimatorBase))
     path = estimator.save_to_disk_for_device(device)
     estimator.set_as_active_calibration_for_device(device)
     logger.debug(