Repository: koaning/scikit-lego (python)

Issue:
[BUG] Can't load a pickled DebugPipeline with a custom log_callback
Hi,

when I try to load a previously fitted and pickled `DebugPipeline` instance that uses a custom `log_callback` from a pickle file I get the following error:

```python
AttributeError: 'Memory' object has no attribute '_'
```

Here's a simple snippet to reproduce the error (tested on Google Colab with `sklego==0.9.4`)
```python
import joblib
import logging
from sklego.pipeline import make_debug_pipeline, default_log_callback
from sklearn.preprocessing import StandardScaler
from sklearn.linear_model import LinearRegression
from sklearn.datasets import make_regression

X, y = make_regression()
pipeline = make_debug_pipeline(StandardScaler(), LinearRegression(), log_callback=default_log_callback)
pipeline.fit(X, y)

joblib.dump(pipeline, './pipeline.pkl')
joblib.load('./pipeline.pkl')
```

Which returns the following stack trace:
```python
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
[<ipython-input-13-7ed979951525>](https://localhost:8080/#) in <cell line: 0>()
     11 
     12 joblib.dump(pipeline, './pipeline.pkl')
---> 13 joblib.load('./pipeline.pkl')

[/usr/local/lib/python3.11/dist-packages/joblib/numpy_pickle.py](https://localhost:8080/#) in load(filename, mmap_mode)
    656                     return load_compatibility(fobj)
    657 
--> 658                 obj = _unpickle(fobj, filename, mmap_mode)
    659     return obj

[/usr/local/lib/python3.11/dist-packages/joblib/numpy_pickle.py](https://localhost:8080/#) in _unpickle(fobj, filename, mmap_mode)
    575     obj = None
    576     try:
--> 577         obj = unpickler.load()
    578         if unpickler.compat_mode:
    579             warnings.warn("The file '%s' has been generated with a "

[/usr/lib/python3.11/pickle.py](https://localhost:8080/#) in load(self)
   1211                     raise EOFError
   1212                 assert isinstance(key, bytes_types)
-> 1213                 dispatch[key[0]](self)
   1214         except _Stop as stopinst:
   1215             return stopinst.value

[/usr/lib/python3.11/pickle.py](https://localhost:8080/#) in load_reduce(self)
   1588         args = stack.pop()
   1589         func = stack[-1]
-> 1590         stack[-1] = func(*args)
   1591     dispatch[REDUCE[0]] = load_reduce
   1592 

AttributeError: 'Memory' object has no attribute '_'
```

Pull request description:
fix: Allow to pickle/unpickle `DebugPipeline`
# Description

Fixes #735 

## Type of change

- [x] Bug fix (non-breaking change which fixes an issue)
- [ ] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)


## Checklist:

- [x] My code follows the style guidelines (ruff)
- [x] I have commented my code, particularly in hard-to-understand areas
- [ ] I have made corresponding changes to the documentation (also to the readme.md)
- [x] I have added tests that prove my fix is effective or that my feature works
- [ ] I have added tests to check whether the new feature adheres to the sklearn convention
- [x] New and existing unit tests pass locally with my changes


Patch:
diff --git a/sklego/pipeline.py b/sklego/pipeline.py
index 9e2efb1..a0dff26 100644
--- a/sklego/pipeline.py
+++ b/sklego/pipeline.py
@@ -309,6 +309,48 @@ class DebugPipeline(Pipeline):
         if self._log_callback == "default":
             self._log_callback = default_log_callback
 
+    def __getstate__(self):
+        """
+        Prepare the state of the object for pickling.
+
+        This method is called when the object is being pickled. It ensures that the `Memory` object is in its original
+        state by temporarily restoring the original `cache` method and removing the custom `_cache` attribute. This is
+        necessary because the custom `_cache` attribute is not picklable and would cause errors during the pickling
+        process.
+
+        Returns
+        -------
+        dict
+            The state of the object to be pickled.
+        """
+        state = self.__dict__.copy()
+        if hasattr(self._memory, "_cache"):
+            self._memory.cache = self._memory._cache
+            del self._memory._cache
+        return state
+
+    def __setstate__(self, state):
+        """
+        Restore the state of the object from the pickled state.
+
+        This method is called when the object is being unpickled. It restores the state of the object and re-applies the
+        custom `_cache` attribute by wrapping the `cache` method with the logging wrapper (`_cache_with_function_log_statement`).
+        This ensures that the `Memory` object has the custom `_cache` attribute after unpickling.
+
+        Parameters
+        ----------
+        state : dict
+            The state of the object to be restored.
+        """
+        self.__dict__.update(state)
+        if self._log_callback is not None:
+            self._memory = check_memory(self._memory)
+            if not hasattr(self._memory, "_cache"):
+                self._memory._cache = self._memory.cache
+            self._memory.cache = _cache_with_function_log_statement(self._log_callback).__get__(
+                self._memory, self._memory.__class__
+            )
+
 
 def make_debug_pipeline(*steps, **kwargs):
     """Construct a `DebugPipeline` from the given estimators.


Test patch:
diff --git a/tests/test_pipeline/test_debug_pipeline.py b/tests/test_pipeline/test_debug_pipeline.py
index e3b1f53..a50f8b7 100644
--- a/tests/test_pipeline/test_debug_pipeline.py
+++ b/tests/test_pipeline/test_debug_pipeline.py
@@ -1,8 +1,11 @@
 import logging
 
+import joblib
 import pytest
 from sklearn import datasets
 from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.datasets import make_regression
+from sklearn.linear_model import LinearRegression
 from sklearn.model_selection import GridSearchCV
 from sklearn.multiclass import (
     OneVsOneClassifier,
@@ -13,7 +16,7 @@ from sklearn.pipeline import FeatureUnion
 from sklearn.preprocessing import StandardScaler
 from sklearn.svm import LinearSVC
 
-from sklego.pipeline import DebugPipeline, make_debug_pipeline
+from sklego.pipeline import DebugPipeline, default_log_callback, make_debug_pipeline
 
 IRIS = datasets.load_iris()
 
@@ -177,3 +180,23 @@ def test_nameless_step_name_in_logs_when_log_callback_is_default(caplog, nameles
     for _, step in pipe.steps[:-1]:
         assert str(step) in caplog.text, f"{step} should be in: {caplog.text}"
         assert caplog.text.count(str(step)) == 1, f"{step} should be once in {caplog.text}"
+
+
+def test_pickling_unpickling_debug_pipeline(caplog, tmp_path):
+    X, y = make_regression()
+    pipe = make_debug_pipeline(StandardScaler(), LinearRegression(), log_callback=default_log_callback)
+    _ = pipe.fit(X, y)
+
+    pickle_file = tmp_path / "pipeline.pkl"
+    joblib.dump(pipe, pickle_file)
+
+    loaded_pipe = joblib.load(pickle_file)
+
+    # Check if the loaded pipeline works as expected
+    assert loaded_pipe is not None
+    assert loaded_pipe.predict(X) is not None
+
+    caplog.clear()
+    with caplog.at_level(logging.INFO):
+        _ = loaded_pipe.fit(X, y)
+    assert caplog.text, f"Log should be none empty: {caplog.text}"