WIP: Add helper for importing optional dependencies

medkit/_import.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+import importlib
+import inspect
+import sys
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    import types
+
+
+__all__ = ["import_or_raise", "import_optional"]
+
+
+def import_or_raise(
+    name: str,
+    package: str | None = None,
+    note: str | None = None,
+) -> types.ModuleType:
+    """Import module or raise an error message.
+
+    Parameters
+    ----------
+    name : str
+        Module to import.
+    package : str, optional
+        Relative package to import the module from.
+    note : str, optional
+        Add this note to the raised exception.
+
+    Returns
+    -------
+    ModuleType
+        Successfully imported module
+
+    Raises
+    ------
+    ModuleNotFoundError
+        In case the requested import failed.
+    """
+    try:
+        module = importlib.import_module(name, package)
+    except ModuleNotFoundError as err:
+        if sys.version_info >= (3, 11):
+            if note:
+                err.add_note(note)
+            raise
+        message = "\n".join([str(err), note or ""])
+        raise ModuleNotFoundError(message) from err
+    return module
+
+
+def import_optional(
+    name: str,
+    package: str | None = None,
+    extra: str | None = None,
+) -> types.ModuleType:
+    """Import an optional dependency or raise an appropriate error message.
+
+    Parameters
+    ----------
+    name : str
+        Module to import.
+    package : str, optional
+        Relative package to import the module from.
+    extra : str, optional
+        Group of optional dependencies to suggest installing if the import fails.
+        If unspecified, assume the extra is named after the caller's module.
+
+    Returns
+    -------
+    ModuleType
+        Successfully imported module
+
+    Raises
+    ------
+    ModuleNotFoundError
+        In case the requested import failed.
+    """
+    if not extra:
+        calling_module = inspect.getmodulename(inspect.stack()[1][1])
+        extra = calling_module.replace("_", "-") if calling_module else None
+    note = (
+        (f"Consider installing the appropriate extra with:\n" f"pip install 'medkit-lib[{extra}]'") if extra else None
+    )
+    return import_or_raise(name=name, package=package, note=note)