Fixup mypy issues

equinor · Nov 18, 2024 · 613fc4d · 613fc4d
1 parent 74c1250
commit 613fc4d
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/src/everest/everest_storage.py b/src/everest/everest_storage.py
@@ -163,9 +163,7 @@ def write_to_experiment(self, experiment: _OptimizerOnlyExperiment):
             for df_key, df in batch_data.existing_dataframes.items():
                 df.write_parquet(experiment.optimizer_mount_point / f"{df_key}.parquet")
 
-    def read_from_experiment(
-        self, experiment: _OptimizerOnlyExperiment
-    ) -> EverestStorageDataFrames:
+    def read_from_experiment(self, experiment: _OptimizerOnlyExperiment) -> None:
         self.initial_values = polars.read_parquet(
             experiment.optimizer_mount_point / "initial_values.parquet"
         )
@@ -515,7 +513,7 @@ def _rename_columns(df: polars.DataFrame):
 
     def _store_gradient_results(
         self, config: EnOptConfig, results: FunctionResults
-    ) -> Tuple[polars.DataFrame, polars.DataFrame]:
+    ) -> Tuple[polars.DataFrame, polars.DataFrame, polars.DataFrame, polars.DataFrame]:
         obj_gradient_per_real = polars.from_pandas(
             results.to_dataframe(config, "evaluations").reset_index()
         )
@@ -672,7 +670,7 @@ def _handle_finished_batch_event(self, event: Event):
         converted_results = tuple(
             convert_to_maximize(result) for result in event.results
         )
-        results = []
+        results: List[FunctionResults | GradientResults] = []
         best_value = -np.inf
         best_results = None
         for item in converted_results: