Fix import_optional_module function to return NoneType

cyclops/evaluate/metrics/experimental/metric_dict.py

@@ -37,6 +37,7 @@
         "torchmetrics.metric",
         attribute="Metric",
         error="ignore",
+        return_nonetype=True,
     )
 
 

cyclops/utils/optional.py

@@ -3,14 +3,16 @@
 import importlib
 import importlib.util
 import warnings
-from typing import Any, Literal, Optional
+from types import ModuleType
+from typing import Literal, Optional, Union
 
 
 def import_optional_module(
     name: str,
     attribute: Optional[str] = None,
     error: Literal["raise", "warn", "ignore"] = "raise",
-) -> Optional[Any]:
+    return_nonetype: bool = False,
+) -> Union[ModuleType, type, None]:
     """Import an optional module.
 
     Parameters
@@ -24,12 +26,34 @@ def import_optional_module(
         - "raise": raise an error if the module cannot be imported.
         - "warn": raise a warning if the module cannot be imported.
         - "ignore": ignore the missing module and return `None`.
+    return_nonetype : bool, optional
+        If `True`, return `NoneType` if the module cannot be imported.
+        Otherwise, return `None`
 
     Returns
     -------
-    Optional[Any]
-        The imported module or attribute from the module, or `None` if the
-        module could not be imported.
+    ModuleType or type or None
+        NoneType or type(None) if the module could not be imported,
+        or the module or attribute if it was imported successfully.
+
+    Raises
+    ------
+    ImportError
+        If the module could not be imported and `error` is set to "raise".
+
+    Warns
+    -----
+    UserWarning
+        If the module could not be imported and `error` is set to "warn".
+
+    Notes
+    -----
+    This function is useful for handling optional dependencies. It will
+    attempt to import the specified module and return it if it is found.
+    If the module is not found, it will raise an ImportError, raise a
+    warning, or return ``None`` or ``NoneType`` based on the value of
+    the `error` parameter. ``NoneType`` is bascially the same as
+    ``type(None)``. This is useful for type hinting.
 
     """
     if error not in ("raise", "warn", "ignore"):
@@ -53,4 +77,7 @@ def import_optional_module(
         if error == "warn":
             warnings.warn(msg, category=ImportWarning, stacklevel=2)
 
+    if return_nonetype:
+        return type(None)
+
     return None

tests/cyclops/utils/test_optional.py

@@ -34,6 +34,13 @@ def test_import_nonexistent_module_ignore():
         error="ignore",
     )
     assert attr is None
+    attr = import_optional_module(
+        "nonexistent_module",
+        attribute="nonexistent_attribute",
+        error="ignore",
+        return_nonetype=True,
+    )
+    assert attr is type(None)
 
 
 def test_import_nonexistent_module_warn():