Add required dependencies error to available_* methods (#1148)

flash/core/model.py

@@ -14,6 +14,7 @@
 import functools
 import inspect
 import pickle
+import re
 from abc import ABCMeta
 from copy import deepcopy
 from importlib import import_module
@@ -267,11 +268,11 @@ def __new__(mcs, *args, **kwargs):
         result = ABCMeta.__new__(mcs, *args, **kwargs)
         if result.required_extras is not None:
             result.__init__ = requires(result.required_extras)(result.__init__)
-            load_from_checkpoint = getattr(result, "load_from_checkpoint", None)
-            if load_from_checkpoint is not None:
-                result.load_from_checkpoint = classmethod(
-                    requires(result.required_extras)(result.load_from_checkpoint.__func__)
-                )
+
+            patterns = ["load_from_checkpoint", "available_*"]  # must match classmethods only
+            regex = "(" + ")|(".join(patterns) + ")"
+            for attribute_name, attribute_value in filter(lambda x: re.match(regex, x[0]), inspect.getmembers(result)):
+                setattr(result, attribute_name, classmethod(requires(result.required_extras)(attribute_value.__func__)))
         return result
 
 

tests/core/test_model.py

@@ -36,6 +36,7 @@
 from flash.core.utilities.imports import (
     _AUDIO_TESTING,
     _GRAPH_TESTING,
+    _IMAGE_AVAILABLE,
     _IMAGE_TESTING,
     _TABULAR_TESTING,
     _TEXT_TESTING,
@@ -317,6 +318,12 @@ class Foo(ImageClassifier):
     assert Foo.available_backbones() == {}
 
 
+@pytest.mark.skipif(_IMAGE_AVAILABLE, reason="image libraries are installed.")
+def test_available_backbones_raises():
+    with pytest.raises(ModuleNotFoundError, match="Required dependencies not available."):
+        _ = ImageClassifier.available_backbones()
+
+
 @ClassificationTask.lr_schedulers
 def custom_steplr_configuration_return_as_instance(optimizer):
     return torch.optim.lr_scheduler.StepLR(optimizer, step_size=10)