ENH Adds InconsistentVersionWarning when unpickling version doesn't match (#25297)

thomasjpfan · glemaitre · web-flow · commit 70c690ebdd48 · 2023-01-12T10:32:23.000+01:00
Co-authored-by: Guillaume Lemaitre &lt;g.lemaitre58@gmail.com&gt;
diff --git a/doc/model_persistence.rst b/doc/model_persistence.rst
@@ -55,6 +55,19 @@ with::
    available `here
    <https://joblib.readthedocs.io/en/latest/persistence.html>`_.
 
+When an estimator is unpickled with a scikit-learn version that is inconsistent
+with the version the estimator was pickled with, a
+:class:`~sklearn.exceptions.InconsistentVersionWarning` is raised. This warning
+can be caught to obtain the original version the estimator was pickled with:
+
+  from sklearn.exceptions import InconsistentVersionWarning
+  warnings.simplefilter("error", InconsistentVersionWarning)
+
+  try:
+      est = pickle.loads("model_from_prevision_version.pickle")
+  except InconsistentVersionWarning as w:
+      print(w.original_sklearn_version)
+
 .. _persistence_limitations:
 
 Security & maintainability limitations
diff --git a/doc/modules/classes.rst b/doc/modules/classes.rst
@@ -443,6 +443,7 @@ Samples generator
    exceptions.DataDimensionalityWarning
    exceptions.EfficiencyWarning
    exceptions.FitFailedWarning
+   exceptions.InconsistentVersionWarning
    exceptions.NotFittedError
    exceptions.UndefinedMetricWarning
 
diff --git a/doc/whats_new/v1.3.rst b/doc/whats_new/v1.3.rst
@@ -100,6 +100,13 @@ Changelog
   out-of-bag scores via the `oob_scores_` or `oob_score_` attributes.
   :pr:`24882` by :user:`Ashwin Mathur <awinml>`.
 
+:mod:`sklearn.exception`
+........................
+- |Feature| Added :class:`exception.InconsistentVersionWarning` which is raised
+  when a scikit-learn estimator is unpickled with a scikit-learn version that is
+  inconsistent with the sckit-learn verion the estimator was pickled with.
+  :pr:`25297` by `Thomas Fan`_.
+
 :mod:`sklearn.feature_extraction`
 .................................
 
diff --git a/sklearn/base.py b/sklearn/base.py
@@ -19,6 +19,7 @@
 from .utils._tags import (
     _DEFAULT_TAGS,
 )
+from .exceptions import InconsistentVersionWarning
 from .utils.validation import check_X_y
 from .utils.validation import check_array
 from .utils.validation import _check_y
@@ -297,15 +298,11 @@ def __setstate__(self, state):
             pickle_version = state.pop("_sklearn_version", "pre-0.18")
             if pickle_version != __version__:
                 warnings.warn(
-                    "Trying to unpickle estimator {0} from version {1} when "
-                    "using version {2}. This might lead to breaking code or "
-                    "invalid results. Use at your own risk. "
-                    "For more info please refer to:\n"
-                    "https://scikit-learn.org/stable/model_persistence.html"
-                    "#security-maintainability-limitations".format(
-                        self.__class__.__name__, pickle_version, __version__
+                    InconsistentVersionWarning(
+                        estimator_name=self.__class__.__name__,
+                        current_sklearn_version=__version__,
+                        original_sklearn_version=pickle_version,
                     ),
-                    UserWarning,
                 )
         try:
             super().__setstate__(state)
diff --git a/sklearn/exceptions.py b/sklearn/exceptions.py
@@ -128,3 +128,38 @@ class PositiveSpectrumWarning(UserWarning):
 
     .. versionadded:: 0.22
     """
+
+
+class InconsistentVersionWarning(UserWarning):
+    """Warning raised when an estimator is unpickled with a inconsistent version.
+
+    Parameters
+    ----------
+    estimator_name : str
+        Estimator name.
+
+    current_sklearn_version : str
+        Current scikit-learn version.
+
+    original_sklearn_version : str
+        Original scikit-learn version.
+    """
+
+    def __init__(
+        self, *, estimator_name, current_sklearn_version, original_sklearn_version
+    ):
+        self.estimator_name = estimator_name
+        self.current_sklearn_version = current_sklearn_version
+        self.original_sklearn_version = original_sklearn_version
+
+    def __str__(self):
+        return (
+            f"Trying to unpickle estimator {self.estimator_name} from version"
+            f" {self.original_sklearn_version} when "
+            f"using version {self.current_sklearn_version}. This might lead to breaking"
+            " code or "
+            "invalid results. Use at your own risk. "
+            "For more info please refer to:\n"
+            "https://scikit-learn.org/stable/model_persistence.html"
+            "#security-maintainability-limitations"
+        )
diff --git a/sklearn/tests/test_base.py b/sklearn/tests/test_base.py
@@ -22,6 +22,7 @@
 from sklearn.tree import DecisionTreeClassifier
 from sklearn.tree import DecisionTreeRegressor
 from sklearn import datasets
+from sklearn.exceptions import InconsistentVersionWarning
 
 from sklearn.base import TransformerMixin
 from sklearn.utils._mocking import MockDataFrame
@@ -397,9 +398,15 @@ def test_pickle_version_warning_is_issued_upon_different_version():
         old_version="something",
         current_version=sklearn.__version__,
     )
-    with pytest.warns(UserWarning, match=message):
+    with pytest.warns(UserWarning, match=message) as warning_record:
         pickle.loads(tree_pickle_other)
 
+    message = warning_record.list[0].message
+    assert isinstance(message, InconsistentVersionWarning)
+    assert message.estimator_name == "TreeBadVersion"
+    assert message.original_sklearn_version == "something"
+    assert message.current_sklearn_version == sklearn.__version__
+
 
 class TreeNoVersion(DecisionTreeClassifier):
     def __getstate__(self):
