From cf2efc077e1b23c4773a84d25c0120a7446071e7 Mon Sep 17 00:00:00 2001 From: "Yngve S. Kristiansen" Date: Thu, 20 Jun 2024 11:50:28 +0200 Subject: [PATCH] (miniwip) refactor it --- src/ert/storage/local_ensemble.py | 2 +- src/ert/storage/realization_state.py | 13 ++++++++----- 2 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/ert/storage/local_ensemble.py b/src/ert/storage/local_ensemble.py index 6633a5b1454..93463c3fab6 100644 --- a/src/ert/storage/local_ensemble.py +++ b/src/ert/storage/local_ensemble.py @@ -159,7 +159,7 @@ def create_realization_dir(realization: int) -> Path: self._realization_dir = create_realization_dir self._realization_states = _MultiRealizationStateDict() - def on_experiment_initialized(self): + def on_experiment_initialized(self) -> None: """ Executes logic that depends on the experiment of the ensemble to exist. For example, if some logic needs to traverse response/parameter configs, diff --git a/src/ert/storage/realization_state.py b/src/ert/storage/realization_state.py index 1fd11ea9bb2..0fe5242ed4c 100644 --- a/src/ert/storage/realization_state.py +++ b/src/ert/storage/realization_state.py @@ -40,7 +40,7 @@ def assign_state(self, src_state: "_SingleRealizationStateDictEntry") -> Self: class _SingleRealizationStateDict: - def __init__(self): + def __init__(self) -> None: self._items_by_kind: Dict[str, Dict[str, _SingleRealizationStateDictEntry]] = {} def _set_item( @@ -147,7 +147,10 @@ def get_parameter(self, key_or_group: str) -> _SingleRealizationStateDictEntry: ) kind_dict = self._lookup_single_kind_dict_for_key(key_or_group) - return kind_dict.get(key_or_group) + + entry = kind_dict.get(key_or_group) + assert entry is not None + return entry def copy(self) -> "_SingleRealizationStateDict": cpy = _SingleRealizationStateDict() @@ -229,16 +232,16 @@ def to_tuples(self) -> List[Tuple[str, str, _SingleRealizationStateDictEntry]]: class _MultiRealizationStateDict: - def __init__(self): + def __init__(self) -> None: self._items: Dict[int, _SingleRealizationStateDict] = {} - def has_response(self, realization: int, key: str): + def has_response(self, realization: int, key: str) -> bool: if realization not in self._items: return False return self._items[realization].has_response_key_or_group(key) - def has_parameter_group(self, realization: int, key: str): + def has_parameter_group(self, realization: int, key: str) -> bool: if realization not in self._items: return False