Add condition to ModuleType (#223)

* Add and test condition

* Fix RTD config

Author: wesselb

.readthedocs.yaml

@@ -17,3 +17,4 @@ python:
 sphinx:
   builder: html
   fail_on_warning: true
+  configuration: docs/conf.py

docs/types.md

@@ -226,6 +226,32 @@ plum.type.ModuleType[tensorflow.python.framework.ops.EagerTensor]
 tensorflow.python.framework.ops.EagerTensor
 ```
 
+You might run into a scenario where an import is only possible when a certain condition
+is satisfied, e.g. a constraint on the package version.
+You can specify a condition with the keyword argument `condition`.
+
+Example:
+
+```python
+>>> def jax_version():
+...     import sys
+...     version_string = sys.modules["jax.version"].__version__
+...     return tuple(int(x) for x in version_string.split("."))
+
+>>> ArrayImpl = Union[
+...     ModuleType(
+...         "jaxlib.xla_extension",
+...         "ArrayImpl",
+...         condition=lambda: jax_version() < (0, 6, 0),
+...     ),
+...     ModuleType(
+...         "jaxlib._jax",
+...         "ArrayImpl",
+...         condition=lambda: jax_version() >= (0, 6, 0),
+...     ),
+... ]
+```
+
 % skip: end
 
 (promisedtype)=

plum/type.py

@@ -2,8 +2,8 @@
 import sys
 import typing
 import warnings
-from collections.abc import Hashable
-from typing import Union, get_args, get_origin
+from collections.abc import Callable, Hashable
+from typing import Optional, Union, get_args, get_origin
 
 from typing_extensions import Self, TypeGuard
 
@@ -91,17 +91,34 @@ class ModuleType(ResolvableType):
         name (str): Name of the type that is promised.
         allow_fail (bool, optional): If the type is does not exist in `module`,
             do not raise an `AttributeError`.
+        condition (Callable[[], bool], optional): A callable that can check a condition,
+            like a package version. This callable will be run whenever `module` has been
+            imported. Only if the callable returns `True`, `name` will be imported
+            from `module`.
     """
 
-    def __init__(self, module: str, name: str, allow_fail: bool = False) -> None:
+    def __init__(
+        self,
+        module: str,
+        name: str,
+        allow_fail: bool = False,
+        condition: Optional[Callable[[], bool]] = None,
+    ) -> None:
         if module in {"__builtin__", "__builtins__"}:
             module = "builtins"
         ResolvableType.__init__(self, f"ModuleType[{module}.{name}]")
         self._name = name
         self._module = module
         self._allow_fail = allow_fail
-
-    def __new__(cls, module: str, name: str, allow_fail: bool = False) -> Self:
+        self._condition = condition
+
+    def __new__(
+        cls,
+        module: str,
+        name: str,
+        allow_fail: bool = False,
+        condition: Optional[Callable[[], bool]] = None,
+    ) -> Self:
         return ResolvableType.__new__(cls, f"ModuleType[{module}.{name}]")
 
     def retrieve(self) -> bool:
@@ -111,6 +128,10 @@ def retrieve(self) -> bool:
             bool: Whether the retrieval succeeded.
         """
         if self._type is None and self._module in sys.modules:
+            # If a condition is given, check the condition before attempting to import.
+            if self._condition is not None and not self._condition():
+                return False
+
             type = sys.modules[self._module]
             for name in self._name.split("."):
                 # If `type` does not contain `name` and `self._allow_fail` is

tests/test_type.py

@@ -86,6 +86,14 @@ def test_moduletype_allow_fail():
     assert not t_allowed.retrieve()
 
 
+def test_moduletype_condition():
+    store = {"condition": False}
+    t = ModuleType("builtins", "int", condition=lambda: store["condition"])
+    assert not t.retrieve()
+    store["condition"] = True
+    assert t.retrieve()
+
+
 def test_is_hint():
     assert not _is_hint(int)
     assert _is_hint(typing.Union[int, float])