Dedupe some state checking logic

src/ert/storage/local_ensemble.py

@@ -225,8 +225,8 @@ def get_realization_mask_with_parameters(self) -> npt.NDArray[np.bool_]:
 
         return np.array(
             [
-                self._parameters_exist_for_realization(i)
-                for i in range(self.ensemble_size)
+                state == RealizationStorageState.INITIALIZED
+                for state in self.get_ensemble_state()
             ]
         )
 
@@ -249,73 +249,11 @@ def get_realization_mask_with_responses(
 
         return np.array(
             [
-                self._responses_exist_for_realization(i, key)
-                for i in range(self.ensemble_size)
+                state == RealizationStorageState.HAS_DATA
+                for state in self.get_ensemble_state()
             ]
         )
 
-    def _parameters_exist_for_realization(self, realization: int) -> bool:
-        """
-        Returns true if all parameters in the experiment have
-        all been saved in the ensemble. If no parameters, return True
-
-        Parameters
-        ----------
-        realization : int
-            Realization index.
-
-        Returns
-        -------
-        exists : bool
-            True if parameters exist for realization.
-        """
-        if not self.experiment.parameter_configuration:
-            return True
-        path = self._realization_dir(realization)
-        return all(
-            (path / (_escape_filename(parameter) + ".nc")).exists()
-            for parameter in self.experiment.parameter_configuration
-        )
-
-    def _responses_exist_for_realization(
-        self, realization: int, key: Optional[str] = None
-    ) -> bool:
-        """
-        Returns true if there are responses in the experiment and they have
-        all been saved in the ensemble
-
-        Parameters
-        ----------
-        realization : int
-            Realization index.
-        key : str, optional
-            Response key to filter realizations. If None, all responses are considered.
-
-        Returns
-        -------
-        exists : bool
-            True if responses exist for realization.
-        """
-
-        if not self.experiment.response_configuration:
-            return True
-        path = self._realization_dir(realization)
-
-        def _has_response(_key: str) -> bool:
-            if _key in self.experiment.response_key_to_response_type:
-                _response_type = self.experiment.response_key_to_response_type[_key]
-                return (path / f"{_response_type}.parquet").exists()
-
-            return (path / f"{_key}.parquet").exists()
-
-        if key:
-            return _has_response(key)
-
-        return all(
-            _has_response(response)
-            for response in self.experiment.response_configuration
-        )
-
     def is_initalized(self) -> List[int]:
         """
         Return the realization numbers where all parameters are internalized. In
@@ -488,14 +426,76 @@ def get_ensemble_state(self) -> List[RealizationStorageState]:
             List of realization states.
         """
 
+        def _parameters_exist_for_realization(realization: int) -> bool:
+            """
+            Returns true if all parameters in the experiment have
+            all been saved in the ensemble. If no parameters, return True
+
+            Parameters
+            ----------
+            realization : int
+                Realization index.
+
+            Returns
+            -------
+            exists : bool
+                True if parameters exist for realization.
+            """
+            if not self.experiment.parameter_configuration:
+                return True
+            path = self._realization_dir(realization)
+            return all(
+                (path / (_escape_filename(parameter) + ".nc")).exists()
+                for parameter in self.experiment.parameter_configuration
+            )
+
+        def _responses_exist_for_realization(
+            realization: int, key: Optional[str] = None
+        ) -> bool:
+            """
+            Returns true if there are responses in the experiment and they have
+            all been saved in the ensemble
+
+            Parameters
+            ----------
+            realization : int
+                Realization index.
+            key : str, optional
+                Response key to filter realizations. If None, all responses are considered.
+
+            Returns
+            -------
+            exists : bool
+                True if responses exist for realization.
+            """
+
+            if not self.experiment.response_configuration:
+                return True
+            path = self._realization_dir(realization)
+
+            def _has_response(_key: str) -> bool:
+                if _key in self.experiment.response_key_to_response_type:
+                    _response_type = self.experiment.response_key_to_response_type[_key]
+                    return (path / f"{_response_type}.parquet").exists()
+
+                return (path / f"{_key}.parquet").exists()
+
+            if key:
+                return _has_response(key)
+
+            return all(
+                _has_response(response)
+                for response in self.experiment.response_configuration
+            )
+
         def _find_state(realization: int) -> RealizationStorageState:
             if self.has_failure(realization):
                 failure = self.get_failure(realization)
                 assert failure
                 return failure.type
-            if self._responses_exist_for_realization(realization):
+            if _responses_exist_for_realization(realization):
                 return RealizationStorageState.HAS_DATA
-            if self._parameters_exist_for_realization(realization):
+            if _parameters_exist_for_realization(realization):
                 return RealizationStorageState.INITIALIZED
             else:
                 return RealizationStorageState.UNDEFINED