From 251d77754b0c5155842db7e6446088456601bcb7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jan=20C=2E=20Riven=C3=A6s?= <jan@rivenaes.net>
Date: Thu, 24 Oct 2024 14:51:42 +0200
Subject: [PATCH] MAINT: resolve import roxar vs import rmsapi

---
 src/xtgeo/__init__.py                     |  7 +--
 src/xtgeo/cube/_cube_roxapi.py            |  6 +--
 src/xtgeo/grid3d/_grid_roxapi.py          | 32 ++++++--------
 src/xtgeo/grid3d/_gridprop_roxapi.py      | 24 ++++++----
 src/xtgeo/roxutils/_roxar_loader.py       | 54 +++++++++++++++++++++++
 src/xtgeo/roxutils/_roxutils_etc.py       |  7 +--
 src/xtgeo/roxutils/roxutils.py            |  6 +--
 src/xtgeo/surface/_regsurf_roxapi.py      | 13 +-----
 src/xtgeo/well/_blockedwell_roxapi.py     |  6 +--
 src/xtgeo/xyz/_xyz_roxapi.py              | 35 +++++----------
 tests/test_roxarapi/test_roxarapi_reek.py | 10 ++---
 11 files changed, 103 insertions(+), 97 deletions(-)
 create mode 100644 src/xtgeo/roxutils/_roxar_loader.py

diff --git a/src/xtgeo/__init__.py b/src/xtgeo/__init__.py
index cf3b99be3..0eb909567 100644
--- a/src/xtgeo/__init__.py
+++ b/src/xtgeo/__init__.py
@@ -35,11 +35,6 @@ def _xprint(msg):
 
 _xprint("XTGEO __init__ ...")
 
-ROXAR = True
-try:
-    import roxar
-except ImportError:
-    ROXAR = False
 
 try:
     from xtgeo.common.version import __version__, version
@@ -159,7 +154,7 @@ def _xprint(msg):
 _xprint("XTGEO __init__ done")
 
 # Remove symbols imported for internal use
-del os, platform, sys, timeit, warnings, TIME0, DEBUG, ROXAR, _timer, _xprint
+del os, platform, sys, timeit, warnings, TIME0, DEBUG, _timer, _xprint
 
 # Let type-checkers know what is exported
 __all__ = [
diff --git a/src/xtgeo/cube/_cube_roxapi.py b/src/xtgeo/cube/_cube_roxapi.py
index d64c63cbc..c98528cbd 100644
--- a/src/xtgeo/cube/_cube_roxapi.py
+++ b/src/xtgeo/cube/_cube_roxapi.py
@@ -18,6 +18,7 @@
 
 from xtgeo.common import XTGeoDialog, null_logger
 from xtgeo.roxutils import RoxUtils
+from xtgeo.roxutils._roxar_loader import roxar
 
 xtg = XTGeoDialog()
 
@@ -132,11 +133,6 @@ def export_cube_roxapi(
 def _roxapi_export_cube(
     self, proj, rox, name, folder=None, domain="time", compression=("wavelet", 5)
 ):  # type: ignore # pragma: no cover
-    try:
-        import roxar  # type: ignore
-    except ImportError:
-        roxar = None
-
     logger.info(
         "There are issues with compression %s, hence it is ignored", compression
     )
diff --git a/src/xtgeo/grid3d/_grid_roxapi.py b/src/xtgeo/grid3d/_grid_roxapi.py
index 46ddfed9c..ac67d7e8f 100644
--- a/src/xtgeo/grid3d/_grid_roxapi.py
+++ b/src/xtgeo/grid3d/_grid_roxapi.py
@@ -11,6 +11,7 @@
 from xtgeo import _cxtgeo
 from xtgeo.common import XTGeoDialog, null_logger
 from xtgeo.common.constants import UNDEF_LIMIT
+from xtgeo.roxutils._roxar_loader import roxar, roxar_grids
 from xtgeo.roxutils.roxutils import RoxUtils
 
 xtg = XTGeoDialog()
@@ -18,12 +19,13 @@
 logger = null_logger(__name__)
 
 if TYPE_CHECKING:
-    import contextlib
-
     from xtgeo.grid3d.grid import Grid
+    # from xtgeo.roxutils._roxar_loader import RoxarGrid3DType
+
+    if roxar is not None:
+        from roxar import grids as RoxarGridType
+        from roxar.grids import Grid3D as RoxarGrid3DType
 
-    with contextlib.suppress(ImportError):
-        import roxar
 
 # self is Grid() instance
 
@@ -47,13 +49,13 @@ def import_grid_roxapi(
     return _import_grid_roxapi(rox, gname, realisation, info)
 
 
-def _display_roxapi_grid_info(roxgrid: roxar.grids.Grid3D) -> None:
+def _display_roxapi_grid_info(roxgrid: RoxarGrid3DType) -> None:
     """Push info to screen (mostly for debugging), experimental."""
 
     indexer = roxgrid.grid_indexer
     ncol, nrow, _ = indexer.dimensions
 
-    xtg.say("ROXAPI with support for CornerPointGeometry")
+    xtg.say("ROXAPI with support for CornerPointGridGeometry")
     geom = roxgrid.get_geometry()
     defined_cells = geom.get_defined_cells()
     xtg.say(f"Defined cells \n{defined_cells}")
@@ -103,7 +105,7 @@ def _import_grid_roxapi(
     return result
 
 
-def _convert_to_xtgeo_grid(roxgrid: roxar.grids.Grid3D, gname: str) -> dict[str, Any]:
+def _convert_to_xtgeo_grid(roxgrid: RoxarGrid3DType, gname: str) -> dict[str, Any]:
     """Convert from roxar CornerPointGeometry to xtgeo, version 2 using _xtgformat=2."""
     indexer = roxgrid.grid_indexer
 
@@ -204,10 +206,6 @@ def _export_grid_cornerpoint_roxapi_v1(
     self: Grid, rox: RoxUtils, gname: str, realisation: int, info: bool
 ) -> None:
     """Convert xtgeo geometry to pillar spec in ROXAPI and store."""
-    try:
-        from roxar.grids import CornerPointGridGeometry as CPG
-    except ImportError:
-        raise RuntimeError("Cannot load Roxar module")
 
     logger.info("Load grid via CornerPointGridGeometry...")
 
@@ -215,7 +213,8 @@ def _export_grid_cornerpoint_roxapi_v1(
     grid_model.set_empty(realisation)
     grid = grid_model.get_grid(realisation)
 
-    geom = CPG.create(self.dimensions)
+    roxar_grid_: RoxarGridType = roxar_grids  # for mypy
+    geom = roxar_grid_.CornerPointGridGeometry.create(self.dimensions)
 
     logger.info(geom)
     scopy = self.copy()
@@ -272,16 +271,13 @@ def _export_grid_cornerpoint_roxapi_v2(
     self: Grid, rox: RoxUtils, gname: str, realisation: int, info: bool
 ) -> None:
     """Convert xtgeo geometry to pillar spec in ROXAPI and store _xtgformat=2."""
-    try:
-        from roxar.grids import CornerPointGridGeometry as CPG
-    except ImportError:
-        raise RuntimeError("Cannot load Roxar module")
 
     grid_model = rox.project.grid_models.create(gname)
     grid_model.set_empty(realisation)
     grid = grid_model.get_grid(realisation)
 
-    geom = CPG.create(self.dimensions)
+    roxar_grids_: RoxarGridType = roxar_grids  # for mypy
+    geom = roxar_grids_.CornerPointGridGeometry.create(self.dimensions)
 
     scopy = self.copy()
     scopy.make_zconsistent()
@@ -332,7 +328,7 @@ def _export_grid_cornerpoint_roxapi_v2(
     del scopy
 
 
-def _set_subgrids(self: Grid, rox: RoxUtils, grid: roxar.grids.Grid3D) -> None:
+def _set_subgrids(self: Grid, rox: RoxUtils, grid: RoxarGrid3DType) -> None:
     """Export the subgrid index (zones) to Roxar API.
 
     From roxar API:
diff --git a/src/xtgeo/grid3d/_gridprop_roxapi.py b/src/xtgeo/grid3d/_gridprop_roxapi.py
index 466668c3a..0c222221d 100644
--- a/src/xtgeo/grid3d/_gridprop_roxapi.py
+++ b/src/xtgeo/grid3d/_gridprop_roxapi.py
@@ -3,7 +3,6 @@
 
 from __future__ import annotations
 
-import contextlib
 from typing import TYPE_CHECKING, Any, Literal
 
 import numpy as np
@@ -12,20 +11,26 @@
 from xtgeo.common import null_logger
 from xtgeo.common.constants import UNDEF, UNDEF_INT, UNDEF_INT_LIMIT, UNDEF_LIMIT
 from xtgeo.roxutils import RoxUtils
-
-with contextlib.suppress(ImportError):
-    import roxar
+from xtgeo.roxutils._roxar_loader import roxar
 
 if TYPE_CHECKING:
     from xtgeo.grid3d.grid_property import GridProperty
 
+    if roxar is not None:
+        from roxar import Project as RoxarProjectType
+        from roxar.grids import Grid3D as RoxarGrid3DType
+
 logger = null_logger(__name__)
 
 VALID_ROXAR_DTYPES = [np.uint8, np.uint16, np.float32]
 
 
 def import_prop_roxapi(
-    project: roxar.Project, gname: str, pname: str, realisation: int, faciescodes: bool
+    project: RoxarProjectType,
+    gname: str,
+    pname: str,
+    realisation: int,
+    faciescodes: bool,
 ) -> dict[str, Any]:
     """Import a Property via ROXAR API spec."""
     logger.info("Opening RMS project ...")
@@ -108,7 +113,7 @@ def _convert_to_xtgeo_prop(
 
 def export_prop_roxapi(
     self: GridProperty,
-    project: roxar.Project,
+    project: RoxarProjectType,
     gname: str,
     pname: str,
     realisation: int = 0,
@@ -133,7 +138,7 @@ def export_prop_roxapi(
 def _store_in_roxar(
     self: GridProperty,
     pname: str,
-    roxgrid: roxar.grid.Grid3D,
+    roxgrid: RoxarGrid3DType,
     realisation: int,
     casting: Literal["no", "equiv", "safe", "same_kind", "unsafe"] | None,
 ) -> None:
@@ -163,10 +168,11 @@ def _store_in_roxar(
 
     pvalues = roxgrid.get_grid(realisation=realisation).generate_values(data_type=dtype)
 
+    roxtype: Any = roxar  # needed for mypy
     roxar_property_type = (
-        roxar.GridPropertyType.discrete
+        roxtype.GridPropertyType.discrete
         if self.isdiscrete
-        else roxar.GridPropertyType.continuous
+        else roxtype.GridPropertyType.continuous
     )
 
     pvalues[cellno] = val3d[iind, jind, kind]
diff --git a/src/xtgeo/roxutils/_roxar_loader.py b/src/xtgeo/roxutils/_roxar_loader.py
new file mode 100644
index 000000000..5f581dfd9
--- /dev/null
+++ b/src/xtgeo/roxutils/_roxar_loader.py
@@ -0,0 +1,54 @@
+"""Loading rmsapi or roxar module based on availability"""
+
+import importlib
+from typing import TYPE_CHECKING, Any, Optional, TypeVar
+
+# Roxar*Type are placeholders for the actual types which are not known at this point,
+# and these may be replaced in modules by e.g:
+# ---------
+#     from xtgeo.roxutils._roxar_loader import roxar
+#     if TYPE_CHECKING:
+#         from xtgeo.grid3d.grid_property import GridProperty
+#         if roxar is not None:
+#            from roxar.grids import Grid3D as RoxarGridType
+# ---------
+
+
+def load_module(module_name: str) -> Optional[Any]:
+    try:
+        return importlib.import_module(module_name)
+    except ImportError:
+        return None
+
+
+rmsapi = load_module("rmsapi")
+roxar = rmsapi if rmsapi else load_module("roxar")
+
+if rmsapi:
+    roxar_well_picks = load_module("rmsapi.well_picks")
+    roxar_jobs = load_module("rmsapi.jobs")
+    roxar_grids = load_module("rmsapi.grids")
+else:
+    roxar_well_picks = load_module("roxar.well_picks")
+    roxar_jobs = load_module("roxar.jobs")
+    roxar_grids = load_module("roxar.grids")
+
+# Use TypeVar correctly
+RoxarType = TypeVar("RoxarType")
+RoxarGridType = TypeVar("RoxarGridType")
+RoxarGrid3DType = TypeVar("RoxarGrid3DType")
+RoxarWellpicksType = TypeVar("RoxarWellpicksType")
+RoxarJobsType = TypeVar("RoxarJobsType")
+
+if TYPE_CHECKING:
+    if roxar is not None:
+        import roxar.grids as RoxarGridType  # noqa  # type: ignore
+        from roxar.grids import Grid3D as RoxarGrid3DType  # noqa  # type: ignore
+        from roxar.jobs import Jobs as RoxarJobsType  # noqa  # type: ignore
+        from roxar.well_picks import WellPicks as RoxarWellpicksType  # noqa  # type: ignore
+
+# Explicitly type the roxar* as Optional to indicate it may be None
+roxar: Optional[Any] = roxar
+roxar_grids: Optional[Any] = roxar_grids
+roxar_well_picks: Optional[Any] = roxar_well_picks
+roxar_jobs: Optional[Any] = roxar_jobs
diff --git a/src/xtgeo/roxutils/_roxutils_etc.py b/src/xtgeo/roxutils/_roxutils_etc.py
index b36e4eed2..d999ab212 100644
--- a/src/xtgeo/roxutils/_roxutils_etc.py
+++ b/src/xtgeo/roxutils/_roxutils_etc.py
@@ -1,12 +1,9 @@
 """Private module for etc functions"""
 
-import contextlib
-
-with contextlib.suppress(ImportError):
-    import roxar
-
 from xtgeo.common.log import null_logger
 
+from ._roxar_loader import roxar
+
 logger = null_logger(__name__)
 
 
diff --git a/src/xtgeo/roxutils/roxutils.py b/src/xtgeo/roxutils/roxutils.py
index 10e02099f..6cc6eecac 100644
--- a/src/xtgeo/roxutils/roxutils.py
+++ b/src/xtgeo/roxutils/roxutils.py
@@ -1,16 +1,12 @@
 """Module for simplifying various operation in the Roxar python interface."""
 
-import contextlib
-
 from packaging.version import parse as versionparse
 
-with contextlib.suppress(ImportError):
-    import roxar
-
 from xtgeo.common.log import null_logger
 from xtgeo.common.xtgeo_dialog import XTGeoDialog
 
 from . import _roxutils_etc
+from ._roxar_loader import roxar
 
 xtg = XTGeoDialog()
 logger = null_logger(__name__)
diff --git a/src/xtgeo/surface/_regsurf_roxapi.py b/src/xtgeo/surface/_regsurf_roxapi.py
index 3ccb08c8c..cb6c745ea 100644
--- a/src/xtgeo/surface/_regsurf_roxapi.py
+++ b/src/xtgeo/surface/_regsurf_roxapi.py
@@ -8,6 +8,7 @@
 
 from xtgeo.common.log import null_logger
 from xtgeo.roxutils import RoxUtils
+from xtgeo.roxutils._roxar_loader import roxar
 
 logger = null_logger(__name__)
 
@@ -214,12 +215,6 @@ def _roxapi_export_surface(
             )
         # here a workound; trends.surfaces are read-only in Roxar API, but is seems
         # that load() in RMS is an (undocumented?) workaround...
-        try:
-            import roxar
-        except ImportError as err:
-            raise ImportError(
-                "roxar not available, this functionality is not available"
-            ) from err
 
         roxsurf = proj.trends.surfaces[name]
         with tempfile.TemporaryDirectory() as tmpdir:
@@ -234,12 +229,6 @@ def _roxapi_export_surface(
 
 def _xtgeo_to_roxapi_grid(self):  # pragma: no cover
     # Create a 2D grid
-    try:
-        import roxar
-    except ImportError as err:
-        raise ImportError(
-            "roxar not available, this functionality is not available"
-        ) from err
 
     return roxar.RegularGrid2D.create(
         x_origin=self.xori,
diff --git a/src/xtgeo/well/_blockedwell_roxapi.py b/src/xtgeo/well/_blockedwell_roxapi.py
index f1cb1641d..375649083 100644
--- a/src/xtgeo/well/_blockedwell_roxapi.py
+++ b/src/xtgeo/well/_blockedwell_roxapi.py
@@ -9,11 +9,7 @@
 from xtgeo.common.exceptions import WellNotFoundError
 from xtgeo.common.log import null_logger
 from xtgeo.roxutils import RoxUtils
-
-try:
-    import roxar  # type: ignore
-except ImportError:
-    roxar = None
+from xtgeo.roxutils._roxar_loader import roxar
 
 logger = null_logger(__name__)
 
diff --git a/src/xtgeo/xyz/_xyz_roxapi.py b/src/xtgeo/xyz/_xyz_roxapi.py
index 3f8a7376e..367552ff5 100644
--- a/src/xtgeo/xyz/_xyz_roxapi.py
+++ b/src/xtgeo/xyz/_xyz_roxapi.py
@@ -11,17 +11,16 @@
 import numpy as np
 import pandas as pd
 
-from xtgeo import ROXAR  # type: ignore
 from xtgeo.common._xyz_enum import _AttrName
 from xtgeo.common.constants import UNDEF, UNDEF_INT, UNDEF_INT_LIMIT, UNDEF_LIMIT
 from xtgeo.common.log import null_logger
 from xtgeo.io._file import FileWrapper
 from xtgeo.roxutils import RoxUtils
+from xtgeo.roxutils._roxar_loader import RoxarType, roxar, roxar_well_picks
 from xtgeo.xyz import _xyz_io, points, polygons
 
-if ROXAR:
-    import roxar
-    import roxar.well_picks as roxwp
+if roxar:
+    roxwp = roxar_well_picks
 
 logger = null_logger(__name__)
 
@@ -127,7 +126,7 @@ def _check_presence_in_project(
 
 
 def import_xyz_roxapi(
-    project: roxar.project,
+    project: RoxarType.project,
     name: str,
     category: str | list[str],
     stype: str = "horizons",
@@ -191,13 +190,6 @@ def _roxapi_import_xyz_viafile(
     routine should be replaced!
     """
 
-    try:
-        import roxar
-    except ImportError as err:
-        raise ImportError(
-            "roxar not available, this functionality is not available"
-        ) from err
-
     rox_xyz = _get_roxitem(
         rox,
         name,
@@ -328,7 +320,7 @@ def _add_attributes_to_dataframe(
 
 def export_xyz_roxapi(
     self: points.Points | polygons.Polygons,
-    project: roxar.Project,
+    project: RoxarType.Project,
     name: str,
     category: str | list[str] | None,
     stype: str,
@@ -379,13 +371,6 @@ def _roxapi_export_xyz_viafile(
 
     logger.warning("Realisation %s not in use", realisation)
 
-    try:
-        import roxar
-    except ImportError as err:
-        raise ImportError(
-            "roxar not available, this functionality is not available"
-        ) from err
-
     is_polygons = isinstance(self, polygons.Polygons)
     roxxyz = _get_roxitem(
         rox, name, category, stype, mode="set", is_polygons=is_polygons
@@ -622,7 +607,7 @@ def _roxapi_import_wellpicks(
 
 
 def _create_dataframe_from_wellpicks(
-    well_picks: list[roxar.well_picks.WellPick],
+    well_picks: list[RoxarType.well_picks.WellPick],
     wp_category: Literal["fault", "horizon"],
     attribute_types: dict[str, str],
 ) -> pd.DataFrame:
@@ -748,8 +733,8 @@ def _roxapi_export_xyz_well_picks(
 
 
 def _get_well_pick_set(
-    rox: RoxUtils, well_pick_set: str, rox_wp_type: roxar.WellPickType
-) -> roxar.well_picks.WellPickSet:
+    rox: RoxUtils, well_pick_set: str, rox_wp_type: RoxarType.WellPickType
+) -> RoxarType.well_picks.WellPickSet:
     """
     Function to retrieve a well pick set object. If the given well pick set
     name is not present, it will be created. Otherwise the current well pick
@@ -768,8 +753,8 @@ def _get_well_pick_set(
 def _get_writeable_well_pick_attributes(
     rox: RoxUtils,
     attribute_types: dict[str, str],
-    rox_wp_type: roxar.WellPickType,
-) -> dict[str, roxar.well_picks.WellPickAttribute]:
+    rox_wp_type: RoxarType.WellPickType,
+) -> dict[str, RoxarType.well_picks.WellPickAttribute]:
     """
     Function to retrive a dictionary of regular and user-defined
     roxar WellPickAttribute's. Only writable attributes are
diff --git a/tests/test_roxarapi/test_roxarapi_reek.py b/tests/test_roxarapi/test_roxarapi_reek.py
index 14e29b57c..893fecf98 100644
--- a/tests/test_roxarapi/test_roxarapi_reek.py
+++ b/tests/test_roxarapi/test_roxarapi_reek.py
@@ -11,7 +11,6 @@
 
 from __future__ import annotations
 
-import contextlib
 import pathlib
 from typing import Any
 
@@ -20,9 +19,7 @@
 
 import xtgeo
 from xtgeo.common import XTGeoDialog, null_logger
-
-with contextlib.suppress(ImportError):
-    import roxar
+from xtgeo.roxutils._roxar_loader import roxar, roxar_jobs, roxar_well_picks
 
 xtg = XTGeoDialog()
 logger = null_logger(__name__)
@@ -94,10 +91,9 @@ def _run_blocked_wells_job(
 
     The log names (Poro, Perm, etc) are here hardcoded in params dictionary.
     """
-    import roxar.jobs
 
     # create the block well job and add it to the grid
-    bw_job = roxar.jobs.Job.create(
+    bw_job = roxar_jobs.Job.create(
         ["Grid models", gmname, "Grid"], "Block Wells", "API_BW_Job"
     )
 
@@ -161,7 +157,7 @@ def _run_blocked_wells_job(
 
 def _add_well_pick_to_project(project: Any, well_pick_data: dict, trajectory: str):
     """Add well picks to the project."""
-    import roxar.well_picks as wp
+    wp = roxar_well_picks
 
     mypicks = [
         wp.WellPick.create(