TYP: Adjust ERT typehints to new import locations

Author: tnatt

src/fmu/dataio/_workflows/case/_config.py

@@ -13,7 +13,7 @@
 from fmu.settings import ProjectFMUDirectory
 
 if TYPE_CHECKING:
-    import ert
+    from ert.runpaths import Runpaths as ErtRunpaths
 
 logger: Final = logging.getLogger(__name__)
 logger.setLevel(logging.CRITICAL)
@@ -53,7 +53,7 @@ def validate(self) -> None:
     @classmethod
     def from_presim_workflow(
         cls,
-        run_paths: ert.Runpaths,
+        run_paths: ErtRunpaths,
         args: argparse.Namespace,
         fmu_dir: ProjectFMUDirectory | None = None,
     ) -> Self:

src/fmu/dataio/_workflows/case/_observations.py

@@ -9,7 +9,7 @@
 import pyarrow as pa
 
 if TYPE_CHECKING:
-    import ert
+    from ert.storage import Ensemble as ErtEnsemble
 
 
 logger: Final = logging.getLogger(__name__)
@@ -46,7 +46,7 @@ def _prepare_observations_dataframe(
 
 
 def get_ert_observations_table(
-    ensemble: ert.Ensemble, obs_type: Literal["rft", "summary", "breakthrough"]
+    ensemble: ErtEnsemble, obs_type: Literal["rft", "summary", "breakthrough"]
 ) -> pa.Table | None:
     """Extract observations from ert storage and process it into an arrow table."""
     logger.info(f"Observation type: {obs_type}")

src/fmu/dataio/_workflows/case/_parameters.py

@@ -13,6 +13,7 @@
 
 if TYPE_CHECKING:
     import polars as pl
+    from ert.storage import Ensemble as ErtEnsemble
 
 
 logger: Final = logging.getLogger(__name__)
@@ -51,7 +52,7 @@ def _resolve_pa_field_type(name: str, pa_type: pa.DataType) -> pa.DataType:
 
 
 def _process_parameters(
-    scalars_df: pl.DataFrame, ensemble: ert.Ensemble
+    scalars_df: pl.DataFrame, ensemble: ErtEnsemble
 ) -> tuple[pa.Table, list[int]]:
     """Process parameters into an Arrow table with metadata."""
     import pyarrow as pa
@@ -96,7 +97,7 @@ def _process_parameters(
     return table, realizations
 
 
-def get_ert_parameters_table(ensemble: ert.Ensemble) -> pa.Table | None:
+def get_ert_parameters_table(ensemble: ErtEnsemble) -> pa.Table | None:
     """Exports Ert parameters as a Parquet file as the ensemble level."""
 
     scalars_df = ensemble.load_scalars()

src/fmu/dataio/_workflows/case/main.py

@@ -9,7 +9,7 @@
 import logging
 import shutil
 from pathlib import Path
-from typing import Final
+from typing import TYPE_CHECKING, Final
 
 import ert
 
@@ -34,6 +34,11 @@
 from ._parameters import get_ert_parameters_table
 from .export_case_metadata import ExportCaseMetadata
 
+if TYPE_CHECKING:
+    from ert.runpaths import Runpaths as ErtRunpaths
+    from ert.storage import Ensemble as ErtEnsemble
+
+
 logger: Final = logging.getLogger(__name__)
 logger.setLevel(logging.CRITICAL)
 
@@ -58,8 +63,8 @@
 
 
 def _get_ensemble_name(
-    ensemble: ert.Ensemble,
-    run_paths: ert.Runpaths,
+    ensemble: ErtEnsemble,
+    run_paths: ErtRunpaths,
     casepath: Path,
 ) -> str:
     """Determine ensemble name from run path.
@@ -77,7 +82,7 @@ def _get_ensemble_name(
 
 
 def _queue_ert_parameters(
-    ensemble: ert.Ensemble,
+    ensemble: ErtEnsemble,
     ensemble_name: str,
     workflow_config: CaseWorkflowConfig,
     sumo_uploader: SumoUploaderInterface,
@@ -107,7 +112,7 @@ def _queue_ert_parameters(
 
 
 def _queue_ert_observations_breakthrough(
-    ensemble: ert.Ensemble,
+    ensemble: ErtEnsemble,
     ensemble_name: str,
     workflow_config: CaseWorkflowConfig,
     sumo_uploader: SumoUploaderInterface,
@@ -139,7 +144,7 @@ def _queue_ert_observations_breakthrough(
 
 
 def _queue_ert_observations_rft(
-    ensemble: ert.Ensemble,
+    ensemble: ErtEnsemble,
     ensemble_name: str,
     workflow_config: CaseWorkflowConfig,
     sumo_uploader: SumoUploaderInterface,
@@ -170,7 +175,7 @@ def _queue_ert_observations_rft(
 
 
 def _queue_ert_observations_summary(
-    ensemble: ert.Ensemble,
+    ensemble: ErtEnsemble,
     ensemble_name: str,
     workflow_config: CaseWorkflowConfig,
     sumo_uploader: SumoUploaderInterface,
@@ -233,8 +238,8 @@ def _queue_stratigraphy_mappings(
 
 
 def _upload_files_to_sumo(
-    ensemble: ert.Ensemble,
-    run_paths: ert.Runpaths,
+    ensemble: ErtEnsemble,
+    run_paths: ErtRunpaths,
     workflow_config: CaseWorkflowConfig,
     sumo_uploader: SumoUploaderInterface,
 ) -> None:
@@ -256,8 +261,8 @@ def _upload_files_to_sumo(
 
 
 def _run_workflow(
-    ensemble: ert.Ensemble,
-    run_paths: ert.Runpaths,
+    ensemble: ErtEnsemble,
+    run_paths: ErtRunpaths,
     workflow_config: CaseWorkflowConfig,
 ) -> None:
     """Main workflow entry point."""
@@ -362,8 +367,8 @@ class WfExportCaseMetadata(ert.ErtScript):
     def run(
         self,
         workflow_args: list[str],
-        ensemble: ert.Ensemble,
-        run_paths: ert.Runpaths,
+        ensemble: ErtEnsemble,
+        run_paths: ErtRunpaths,
     ) -> None:
         """Parse arguments and run the workflow."""
         parser = get_parser()
@@ -376,7 +381,7 @@ def run(
 
 
 @ert.plugin(name="fmu_dataio")
-def ertscript_workflow(config: ert.CaseWorkflowConfigs) -> None:
+def ertscript_workflow(config: ert.WorkflowConfigs) -> None:
     """Hook the WfExportCaseMetadata class with documentation into ERT."""
     config.add_workflow(
         WfExportCaseMetadata,

tests/test_ert_integration/test_wf_create_case_metadata.py

@@ -53,6 +53,7 @@
 )
 
 if TYPE_CHECKING:
+    from ert.storage import Ensemble as ErtEnsemble
     from fmu.datamodels.fmu_results.global_configuration import GlobalConfiguration
 
 
@@ -462,7 +463,7 @@ def test_create_case_metadata_collects_ert_parameters_as_expected(
     scalars_and_config = []
 
     def capture_params(
-        ensemble: ert.Ensemble,
+        ensemble: ErtEnsemble,
         ensemble_name: str,
         workflow_config: CaseWorkflowConfig,
         sumo_uploader: SumoUploaderInterface,
@@ -797,7 +798,7 @@ def test_create_case_metadata_collects_rft_observations_as_expected(
     captured_tables = {}
 
     def capture_observation_tables(
-        ensemble: ert.Ensemble,
+        ensemble: ErtEnsemble,
         obs_type: str,
     ) -> None:
         """Captures observation tables from Ert run.
@@ -808,7 +809,7 @@ def capture_observation_tables(
         return df
 
     def mock_create_observation_dataframes(
-        observations: ert.Ensemble,
+        observations: ErtEnsemble,
         shape_registry: ShapeRegistry,
     ) -> dict[str, pl.DataFrame]:
         """mock"""
@@ -868,7 +869,7 @@ def test_create_case_metadata_with_no_observations(
     captured_tables = {}
 
     def capture_observation_tables(
-        ensemble: ert.Ensemble,
+        ensemble: ErtEnsemble,
         obs_type: str,
     ) -> None:
         """Captures rft observations from Ert run"""

tests/test_units/test_workflows/test_case_workflow.py

@@ -22,7 +22,7 @@
 
 @pytest.fixture
 def mock_ensemble() -> Callable[[int], MagicMock]:
-    """Mocks an ert.Ensemble object."""
+    """Mocks an ert.storage.Ensemble object."""
 
     def _mock_ensemble(iteration: int = 0) -> MagicMock:
         """Creates the mocked object."""
@@ -35,7 +35,7 @@ def _mock_ensemble(iteration: int = 0) -> MagicMock:
 
 @pytest.fixture
 def mock_run_paths() -> Callable[[str], MagicMock]:
-    """Mocks and ert.Runpaths object."""
+    """Mocks and ert.runpaths.Runpaths object."""
 
     def _mock_run_paths(runpath: str = "/tmp/realization-0/iter-0") -> MagicMock:
         """Creates the mocked object."""