Add Matplotlib ImportError when using compare_models() (#53)

* raise ImportError

* Update matplotlib.py

- added note to explain this is min.

* parsing version + add test

* choose option 1 + fix test

* add decorator and change test

* Add docstring

Author: gonlairo

cebra/helper.py

@@ -17,6 +17,7 @@
 import urllib
 import warnings
 import zipfile
+from functools import wraps
 from typing import List, Union
 
 import numpy as np
@@ -121,3 +122,40 @@ def get_loader_options(dataset: "cebra.data.Dataset") -> List[str]:
         "The 'get_loader_options' function has been moved to 'cebra.data.helpers' module. "
         "Please update your imports.", DeprecationWarning)
     return cebra.data.helper.get_loader_options
+
+
+def requires_package_version(module, version: str):
+    """Decorator to require a minimum version of a package.
+
+    Args:
+        module: Module to be checked.
+        version: The minimum required version for the module.
+
+    Raises:
+        ImportError: If the specified ``module`` version is less than
+            the required ``version``.
+    """
+
+    required_version = pkg_resources.parse_version(version)
+
+    def _requires_package_version(function):
+
+        @wraps(function)
+        def wrapper(*args, patched_version=None, **kwargs):
+            if patched_version != None:
+                installed_version = pkg_resources.parse_version(
+                    patched_version)  # Use the patched version if provided
+            else:
+                installed_version = pkg_resources.parse_version(
+                    module.__version__)
+
+            if installed_version < required_version:
+                raise ImportError(
+                    f"The function '{function.__name__}' requires {module.__name__} "
+                    f"version {required_version} or higher, but you have {installed_version}. "
+                    f"Please upgrade {module.__name__}.")
+            return function(*args, **kwargs)
+
+        return wrapper
+
+    return _requires_package_version

cebra/integrations/matplotlib.py

@@ -1132,6 +1132,10 @@ def plot_consistency(
     ).plot(**kwargs)
 
 
+from cebra.helper import requires_package_version
+
+
+@requires_package_version(matplotlib, "3.6")
 def compare_models(
     models: List[CEBRA],
     labels: Optional[List[str]] = None,
@@ -1185,8 +1189,10 @@ def compare_models(
         The axis of the generated plot. If no ``ax`` argument was specified, it will be created
         by the function and returned here.
     """
+
     if not isinstance(models, list):
         raise ValueError(f"Invalid list of models, got {type(models)}.")
+
     for model in models:
         if not isinstance(model, CEBRA):
             raise ValueError(

tests/test_plot.py

@@ -14,6 +14,7 @@
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
+import pkg_resources
 import pytest
 import torch
 from sklearn.exceptions import NotFittedError
@@ -143,10 +144,32 @@ def test_plot_loss():
     plt.close()
 
 
+@pytest.mark.parametrize("matplotlib_version",
+                         ["3.3", "3.4.2", "3.5", "3.6", "3.7"])
+def test_compare_models_with_different_versions(matplotlib_version):
+    # example dataset
+    X = np.random.uniform(0, 1, (1000, 2))
+    n_models = 2
+
+    fitted_models = []
+    for _ in range(n_models):
+        fitted_models.append(
+            cebra_sklearn_cebra.CEBRA(max_iterations=10, batch_size=128).fit(X))
+
+    # minimum version of matplotlib
+    minimum_version = "3.6"
+
+    if pkg_resources.parse_version(
+            matplotlib_version) < pkg_resources.parse_version(minimum_version):
+        with pytest.raises(ImportError):
+            cebra_plot.compare_models(models=fitted_models,
+                                      patched_version=matplotlib_version)
+
+
 def test_compare_models():
     # example dataset
-    X = np.random.uniform(0, 1, (1000, 50))
-    n_models = 10
+    X = np.random.uniform(0, 1, (100, 5))
+    n_models = 4
 
     fig = plt.figure(figsize=(5, 5))
     ax = fig.add_subplot()