Refactor snapshots

This commit removes some of the excess classes (SnapshotBuilder,
SnapshotDict) from snapshot.py. It also changes the pydantic classes to
be TypedDicts instead.

src/ert/cli/monitor.py

@@ -106,7 +106,7 @@ def _print_job_errors(self) -> None:
         failed_jobs: Dict[Optional[str], int] = {}
         for snapshot in self._snapshots.values():
             for real in snapshot.reals.values():
-                for job in real.forward_models.values():
+                for job in real["forward_models"].values():
                     if job.get(ids.STATUS) == FORWARD_MODEL_STATE_FAILURE:
                         err = job.get(ids.ERROR)
                         result = failed_jobs.get(err, 0)
@@ -120,7 +120,7 @@ def _get_legends(self) -> str:
         total_count = len(latest_snapshot.reals)
         aggregate = latest_snapshot.aggregate_real_states()
         for state_ in ALL_REALIZATION_STATES:
-            count = aggregate.get(state_, 0)
+            count = aggregate[state_]
             _countstring = f"{count}/{total_count}"
             out = (
                 f"{self._colorize(self.dot, color=REAL_STATE_TO_COLOR[state_])}"

src/ert/ensemble_evaluator/__init__.py

@@ -5,7 +5,7 @@
 from .evaluator import EnsembleEvaluator
 from .event import EndEvent, FullSnapshotEvent, SnapshotUpdateEvent
 from .monitor import Monitor
-from .snapshot import ForwardModel, RealizationSnapshot, Snapshot, SnapshotDict
+from .snapshot import ForwardModel, RealizationSnapshot, Snapshot
 
 __all__ = [
     "EndEvent",
@@ -19,7 +19,6 @@
     "Realization",
     "RealizationSnapshot",
     "Snapshot",
-    "SnapshotDict",
     "SnapshotUpdateEvent",
     "wait_for_evaluator",
 ]

src/ert/ensemble_evaluator/_ensemble.py

@@ -25,7 +25,11 @@
 
 from ._wait_for_evaluator import wait_for_evaluator
 from .config import EvaluatorServerConfig
-from .snapshot import ForwardModel, RealizationSnapshot, Snapshot, SnapshotDict
+from .snapshot import (
+    ForwardModel,
+    RealizationSnapshot,
+    Snapshot,
+)
 from .state import (
     ENSEMBLE_STATE_CANCELLED,
     ENSEMBLE_STATE_FAILED,
@@ -120,25 +124,20 @@ def active_reals(self) -> Sequence[Realization]:
         return list(filter(lambda real: real.active, self.reals))
 
     def _create_snapshot(self) -> Snapshot:
-        reals: Dict[str, RealizationSnapshot] = {}
+        snapshot = Snapshot()
+        snapshot._ensemble_state = ENSEMBLE_STATE_UNKNOWN
         for real in self.active_reals:
-            reals[str(real.iens)] = RealizationSnapshot(
-                active=True,
-                status=REALIZATION_STATE_WAITING,
+            realization = RealizationSnapshot(
+                active=True, status=REALIZATION_STATE_WAITING, forward_models={}
             )
             for index, forward_model in enumerate(real.forward_models):
-                reals[str(real.iens)].forward_models[str(index)] = ForwardModel(
+                realization["forward_models"][str(index)] = ForwardModel(
                     status=FORWARD_MODEL_STATE_START,
                     index=str(index),
                     name=forward_model.name,
                 )
-        top = SnapshotDict(
-            reals=reals,
-            status=ENSEMBLE_STATE_UNKNOWN,
-            metadata=self.metadata,
-        )
-
-        return Snapshot.from_nested_dict(top.model_dump())
+            snapshot.add_realization(str(real.iens), realization)
+        return snapshot
 
     def get_successful_realizations(self) -> List[int]:
         return self.snapshot.get_successful_realizations()