Improve conditional imports

.gitignore

@@ -28,3 +28,5 @@ pyoptsparse/pyNLPQLP/source
 *.pdb
 
 *.pyd
+
+.DS_Store

pyoptsparse/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "2.14.2"
+__version__ = "2.14.3"
 
 from .pyOpt_history import History
 from .pyOpt_variable import Variable

pyoptsparse/pyCONMIN/pyCONMIN.py

@@ -13,11 +13,11 @@
 
 # Local modules
 from ..pyOpt_optimizer import Optimizer
-from ..pyOpt_utils import try_import_compiled_module_from_path
+from ..pyOpt_utils import import_module
 
 # import the compiled module
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
-conmin = try_import_compiled_module_from_path("conmin", THIS_DIR, raise_warning=True)
+conmin = import_module("conmin", [THIS_DIR])
 
 
 class CONMIN(Optimizer):
@@ -30,8 +30,8 @@ def __init__(self, raiseError=True, options={}):
         category = "Local Optimizer"
         defOpts = self._getDefaultOptions()
         informs = self._getInforms()
-        if isinstance(conmin, str) and raiseError:
-            raise ImportError(conmin)
+        if isinstance(conmin, Exception) and raiseError:
+            raise conmin
 
         self.set_options = []
         super().__init__(name, category, defaultOptions=defOpts, informs=informs, options=options)

pyoptsparse/pyIPOPT/pyIPOPT.py

@@ -10,16 +10,12 @@
 # External modules
 import numpy as np
 
-try:
-    # External modules
-    import cyipopt
-except ImportError:
-    cyipopt = None
-
 # Local modules
 from ..pyOpt_optimizer import Optimizer
 from ..pyOpt_solution import SolutionInform
-from ..pyOpt_utils import ICOL, INFINITY, IROW, convertToCOO, extractRows, scaleRows
+from ..pyOpt_utils import ICOL, INFINITY, IROW, convertToCOO, extractRows, import_module, scaleRows
+
+cyipopt = import_module("cyipopt")
 
 
 class IPOPT(Optimizer):
@@ -36,8 +32,8 @@ def __init__(self, raiseError=True, options={}):
         category = "Local Optimizer"
         defOpts = self._getDefaultOptions()
         informs = self._getInforms()
-        if cyipopt is None and raiseError:
-            raise ImportError("Could not import cyipopt")
+        if isinstance(cyipopt, Exception) and raiseError:
+            raise cyipopt
 
         super().__init__(
             name,

pyoptsparse/pyNLPQLP/pyNLPQLP.py

@@ -14,11 +14,11 @@
 # Local modules
 from ..pyOpt_optimizer import Optimizer
 from ..pyOpt_solution import SolutionInform
-from ..pyOpt_utils import try_import_compiled_module_from_path
+from ..pyOpt_utils import import_module
 
 # import the compiled module
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
-nlpqlp = try_import_compiled_module_from_path("nlpqlp", THIS_DIR)
+nlpqlp = import_module("nlpqlp", [THIS_DIR])
 
 
 class NLPQLP(Optimizer):
@@ -31,8 +31,8 @@ def __init__(self, raiseError=True, options={}):
         category = "Local Optimizer"
         defOpts = self._getDefaultOptions()
         informs = self._getInforms()
-        if isinstance(nlpqlp, str) and raiseError:
-            raise ImportError(nlpqlp)
+        if isinstance(nlpqlp, Exception) and raiseError:
+            raise nlpqlp
 
         super().__init__(name, category, defaultOptions=defOpts, informs=informs, options=options)
         # NLPQLP needs Jacobians in dense format

pyoptsparse/pyNSGA2/pyNSGA2.py

@@ -13,11 +13,11 @@
 
 # Local modules
 from ..pyOpt_optimizer import Optimizer
-from ..pyOpt_utils import try_import_compiled_module_from_path
+from ..pyOpt_utils import import_module
 
 # import the compiled module
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
-nsga2 = try_import_compiled_module_from_path("nsga2", THIS_DIR, raise_warning=True)
+nsga2 = import_module("nsga2", [THIS_DIR])
 
 
 class NSGA2(Optimizer):
@@ -32,8 +32,8 @@ def __init__(self, raiseError=True, options={}):
         informs = self._getInforms()
         super().__init__(name, category, defaultOptions=defOpts, informs=informs, options=options)
 
-        if isinstance(nsga2, str) and raiseError:
-            raise ImportError(nsga2)
+        if isinstance(nsga2, Exception) and raiseError:
+            raise nsga2
 
         if self.getOption("PopSize") % 4 != 0:
             raise ValueError("Option 'PopSize' must be a multiple of 4")

pyoptsparse/pyOpt_utils.py

@@ -10,11 +10,12 @@
 """
 
 # Standard Python modules
+import contextlib
 import importlib
 import os
 import sys
 import types
-from typing import Optional, Tuple, Union
+from typing import Literal, Sequence, Tuple, Union
 import warnings
 
 # External modules
@@ -361,9 +362,9 @@ def convertToCSC(mat: Union[dict, spmatrix, ndarray]) -> dict:
 
 def convertToDense(mat: Union[dict, spmatrix, ndarray]) -> ndarray:
     """
-    Take a pyopsparse sparse matrix definition and convert back to a dense
+    Take a pyoptsparse sparse matrix definition and convert back to a dense
     format. This is typically the final step for optimizers with dense constraint
-    jacibians.
+    jacobians.
 
     Parameters
     ----------
@@ -576,40 +577,53 @@ def _broadcast_to_array(name: str, value: ArrayType, n_values: int, allow_none:
     return value
 
 
-def try_import_compiled_module_from_path(
-    module_name: str, path: Optional[str] = None, raise_warning: bool = False
-) -> Union[types.ModuleType, str]:
+@contextlib.contextmanager
+def _prepend_path(path: Union[str, Sequence[str]]):
+    """Context manager which temporarily prepends to `sys.path`."""
+    if isinstance(path, str):
+        path = [path]
+    orig_path = sys.path
+    if path:
+        path = [os.path.abspath(os.path.expandvars(os.path.expanduser(p))) for p in path]
+        sys.path = path + sys.path
+    yield
+    sys.path = orig_path
+    return
+
+
+def import_module(
+    module_name: str,
+    path: Union[str, Sequence[str]] = (),
+    on_error: Literal["raise", "return"] = "return",
+) -> Union[types.ModuleType, Exception]:
     """
     Attempt to import a module from a given path.
 
     Parameters
     ----------
     module_name : str
-        The name of the module
-    path : Optional[str]
-        The path to import from. If None, the default ``sys.path`` is used.
-    raise_warning : bool
-        If true, raise an import warning. By default false.
+        The name of the module.
+    path : Union[str, Sequence[str]]
+        The search path, which will be prepended to ``sys.path``. May be a string, or a sequence of strings.
+    on_error : str
+        Specify behavior when import fails. If "raise", any exception raised during the import will be raised.
+        If "return", any exception during the import will be returned.
 
     Returns
     -------
     Union[types.ModuleType, str]
         If importable, the imported module is returned.
-        If not importable, the error message is instead returned.
+        If not importable, the exception is returned.
     """
-    orig_path = sys.path
-    if path is not None:
-        path = os.path.abspath(os.path.expandvars(os.path.expanduser(path)))
-        sys.path = [path]
-    try:
-        module = importlib.import_module(module_name)
-    except ImportError as e:
-        if raise_warning:
-            warnings.warn(
-                f"{module_name} module could not be imported from {path}.",
-                stacklevel=2,
-            )
-        module = str(e)
-    finally:
-        sys.path = orig_path
+    if on_error.lower() not in ("raise", "return"):
+        raise ValueError("`on_error` must be 'raise' or 'return'.")
+
+    with _prepend_path(path):
+        try:
+            module = importlib.import_module(module_name)
+        except ImportError as e:
+            if on_error.lower() == "raise":
+                raise e
+            else:
+                module = e
     return module

pyoptsparse/pyPSQP/pyPSQP.py

@@ -13,11 +13,11 @@
 # Local modules
 from ..pyOpt_optimizer import Optimizer
 from ..pyOpt_solution import SolutionInform
-from ..pyOpt_utils import try_import_compiled_module_from_path
+from ..pyOpt_utils import import_module
 
 # import the compiled module
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
-psqp = try_import_compiled_module_from_path("psqp", THIS_DIR)
+psqp = import_module("psqp", [THIS_DIR])
 
 
 class PSQP(Optimizer):
@@ -31,8 +31,8 @@ def __init__(self, raiseError=True, options={}):
         defOpts = self._getDefaultOptions()
         informs = self._getInforms()
 
-        if isinstance(psqp, str) and raiseError:
-            raise ImportError(psqp)
+        if isinstance(psqp, Exception) and raiseError:
+            raise psqp
 
         super().__init__(name, category, defaultOptions=defOpts, informs=informs, options=options)
 

pyoptsparse/pyParOpt/ParOpt.py

@@ -4,20 +4,25 @@
 try:
     # External modules
     from paropt.paropt_pyoptsparse import ParOptSparse as ParOpt
-except ImportError:
+except ImportError as e:
 
-    class ParOpt(Optimizer):
-        def __init__(self, raiseError=True, options={}):
-            name = "ParOpt"
-            category = "Local Optimizer"
-            self.defOpts = {}
-            self.informs = {}
-            super().__init__(
-                name,
-                category,
-                defaultOptions=self.defOpts,
-                informs=self.informs,
-                options=options,
-            )
-            if raiseError:
-                raise ImportError("There was an error importing ParOpt")
+    def make_cls(e):
+        class ParOpt(Optimizer):
+            def __init__(self, raiseError=True, options={}):
+                name = "ParOpt"
+                category = "Local Optimizer"
+                self.defOpts = {}
+                self.informs = {}
+                super().__init__(
+                    name,
+                    category,
+                    defaultOptions=self.defOpts,
+                    informs=self.informs,
+                    options=options,
+                )
+                if raiseError:
+                    raise e
+
+        return ParOpt
+
+    ParOpt = make_cls(e)

pyoptsparse/pySLSQP/pySLSQP.py

@@ -15,11 +15,11 @@
 from ..pyOpt_error import pyOptSparseWarning
 from ..pyOpt_optimizer import Optimizer
 from ..pyOpt_solution import SolutionInform
-from ..pyOpt_utils import try_import_compiled_module_from_path
+from ..pyOpt_utils import import_module
 
 # import the compiled module
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
-slsqp = try_import_compiled_module_from_path("slsqp", THIS_DIR, raise_warning=True)
+slsqp = import_module("slsqp", [THIS_DIR])
 
 
 class SLSQP(Optimizer):
@@ -32,8 +32,8 @@ def __init__(self, raiseError=True, options={}):
         category = "Local Optimizer"
         defOpts = self._getDefaultOptions()
         informs = self._getInforms()
-        if isinstance(slsqp, str) and raiseError:
-            raise ImportError(slsqp)
+        if isinstance(slsqp, Exception) and raiseError:
+            raise slsqp
 
         self.set_options = []
         super().__init__(name, category, defaultOptions=defOpts, informs=informs, options=options)

pyoptsparse/pySNOPT/pySNOPT.py

@@ -27,15 +27,15 @@
     INFINITY,
     IROW,
     extractRows,
+    import_module,
     mapToCSC,
     scaleRows,
-    try_import_compiled_module_from_path,
 )
 
 # import the compiled module
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
-_IMPORT_SNOPT_FROM = os.environ.get("PYOPTSPARSE_IMPORT_SNOPT_FROM", THIS_DIR)
-snopt = try_import_compiled_module_from_path("snopt", _IMPORT_SNOPT_FROM)
+_IMPORT_SNOPT_FROM = os.environ.get("PYOPTSPARSE_IMPORT_SNOPT_FROM", None) or THIS_DIR
+snopt = import_module("snopt", [_IMPORT_SNOPT_FROM])
 
 
 class SNOPT(Optimizer):
@@ -68,9 +68,9 @@ def __init__(self, raiseError=True, options: Dict = {}):
 
         informs = self._getInforms()
 
-        if isinstance(snopt, str):
+        if isinstance(snopt, Exception):
             if raiseError:
-                raise ImportError(snopt)
+                raise snopt
             else:
                 version = None
         else:

tests/test_other.py

@@ -6,7 +6,7 @@
 # First party modules
 from pyoptsparse import Optimizers, list_optimizers
 from pyoptsparse.pyOpt_solution import SolutionInform
-from pyoptsparse.pyOpt_utils import try_import_compiled_module_from_path
+from pyoptsparse.pyOpt_utils import import_module
 
 # we have to unset this environment variable because otherwise
 # the snopt module gets automatically imported, thus failing the import test below
@@ -19,13 +19,24 @@ def test_nonexistent_path(self):
         for key in list(sys.modules.keys()):
             if "snopt" in key:
                 sys.modules.pop(key)
-        with self.assertWarns(UserWarning):
-            module = try_import_compiled_module_from_path("snopt", "/a/nonexistent/path", raise_warning=True)
-            self.assertTrue(isinstance(module, str))
+        with self.assertRaises(ImportError):
+            import_module("snopt", ["/a/nonexistent/path"], on_error="raise")
+
+    def test_import_standard(self):
+        loaded = import_module("os")
+        assert loaded.__name__ == "os"
+
+    def test_import_nonexistent(self):
+        with self.assertRaises(ImportError):
+            _ = import_module("not_a_module", on_error="raise")
+
+        e = import_module("not_a_module", on_error="return")
+        assert isinstance(e, Exception)
+        assert "No module" in str(e)
 
     def test_sys_path_unchanged(self):
         path = tuple(sys.path)
-        try_import_compiled_module_from_path("snopt", "/some/path")
+        import_module("somemodule", ["/some/path"])
         self.assertEqual(tuple(sys.path), path)