Convert pre-treatment references to class instances

Author: scholtzan

pensieve/statistics.py

@@ -107,6 +107,22 @@ def transform(self, df: DataFrame, metric: str) -> "StatisticResultCollection":
     @classmethod
     def from_dict(cls, config_dict: Dict[str, Any]):
         """Create a class instance with the specified config parameters."""
+        if "pre_treatments" in config_dict:
+            # convert pre-treatments to class instances
+            pre_treatments = []
+
+            for pre_treatment_name in config_dict["pre_treatments"]:
+                found = False
+                for pre_treatment in PreTreatment.__subclasses__():
+                    if pre_treatment.name() == pre_treatment_name:
+                        pre_treatments.append(pre_treatment)
+                        found = True
+
+                if not found:
+                    raise ValueError(f"Could not find pre-treatment {pre_treatment_name}")
+
+            config_dict["pre_treatments"] = pre_treatments
+
         return cls(**config_dict)
 
 

pensieve/tests/integration/test_analysis_integration.py

@@ -77,7 +77,7 @@ def setup(self, client):
             client.query(source_file.read()).result()
             yield
 
-            # client.delete_dataset(self.test_dataset, delete_contents=True, not_found_ok=True)
+            client.delete_dataset(self.test_dataset, delete_contents=True, not_found_ok=True)
 
     def test_metrics(self, client):
         experiment = Experiment(