diff --git a/frappe/model/base_document.py b/frappe/model/base_document.py index f02f9ac23e..1dbe0f1b0b 100644 --- a/frappe/model/base_document.py +++ b/frappe/model/base_document.py @@ -70,6 +70,25 @@ UNPICKLABLE_KEYS = ( ) +def _reconstruct_extended_instance(doctype, state): + """Helper function to reconstruct an extended class instance during unpickling. + + This function is called during unpickling to recreate the extended class + based on current hooks and restore the instance state. + """ + # Get the current extended class (uses caching from get_controller) + extended_class = get_controller(doctype) + instance = extended_class.__new__(extended_class) + + # Use __setstate__ if available, otherwise directly update __dict__ + if hasattr(instance, "__setstate__"): + instance.__setstate__(state) + else: + instance.__dict__.update(state) + + return instance + + def get_controller(doctype): """Return the locally cached **class** object of the given DocType. @@ -160,7 +179,24 @@ def get_extended_class(base_class, doctype): # Create the extended class by combining extension classes with base class # Extension classes come first in MRO, then base class class_name = f"Extended{base_class.__name__}" - extended_class = type(class_name, (*extension_classes, base_class), {}) + + def __reduce__(self): + """Make extended class instances pickle-able. + + When unpickling, this will use get_controller() to recreate the extended class + based on current hooks, ensuring the instance respects the current environment. + Respects the BaseDocument's __getstate__ method for proper state handling. + """ + + return (_reconstruct_extended_instance, (self.doctype, self.__getstate__())) + + extended_class = type( + class_name, + (*extension_classes, base_class), + { + "__reduce__": __reduce__, + }, + ) return extended_class diff --git a/frappe/tests/test_base_document.py b/frappe/tests/test_base_document.py index e08e9def5b..7f22be0be6 100644 --- a/frappe/tests/test_base_document.py +++ b/frappe/tests/test_base_document.py @@ -1,3 +1,5 @@ +import pickle + import frappe from frappe.desk.doctype.todo.todo import ToDo from frappe.model.base_document import BaseDocument, get_extended_class @@ -141,3 +143,65 @@ class TestBaseDocument(IntegrationTestCase): # Check that the error message mentions the invalid path error_message = str(context.exception) self.assertIn(path_to_invalid_extension, error_message) + + def test_extended_class_is_pickleable(self): + """Test that extended class instances can be pickled and unpickled correctly""" + from frappe.desk.doctype.todo.todo import ToDo + + # Mock the hooks to include extensions + extensions = ["frappe.tests.test_base_document.TestToDoExtension"] + + with self.patch_hooks({"extend_doctype_class": {"ToDo": extensions}}): + extended_class = get_extended_class(ToDo, "ToDo") + + # Create an instance with some data + original_instance = extended_class( + {"doctype": "ToDo", "description": "Test ToDo for pickling", "status": "Open"} + ) + + # Set a custom attribute from extension + original_instance.validate() # This sets custom_validation_called = True + original_instance.custom_attribute = "test_value" + + # Test that __getstate__ properly excludes unpicklable values + state = original_instance.__getstate__() + # These should be excluded by BaseDocument's __getstate__ + for unpicklable_key in ["meta", "permitted_fieldnames", "_weakref"]: + self.assertNotIn(unpicklable_key, state) + + # Pickle the instance + pickled_data = pickle.dumps(original_instance) + + # Clear the controller cache to ensure we're not using cached classes + clear_todo_controller_cache() + + try: + # Unpickle the instance (this should recreate the extended class) + unpickled_instance = pickle.loads(pickled_data) + finally: + # Always clean up the controller cache to prevent test pollution + clear_todo_controller_cache() + + # Test that the unpickled instance is of the extended class type + self.assertEqual(unpickled_instance.__class__.__name__, f"Extended{ToDo.__name__}") + + # Test that the instance data is preserved + self.assertEqual(unpickled_instance.doctype, "ToDo") + self.assertEqual(unpickled_instance.description, "Test ToDo for pickling") + self.assertEqual(unpickled_instance.status, "Open") + self.assertEqual(unpickled_instance.custom_attribute, "test_value") + self.assertTrue(getattr(unpickled_instance, "custom_validation_called", False)) + + # Test that extension methods are still available + self.assertTrue(hasattr(unpickled_instance, "extension_method")) + self.assertEqual(unpickled_instance.extension_method(), "extension_method_called") + + # Test that original ToDo methods are still available + self.assertTrue(hasattr(unpickled_instance, "on_update")) + self.assertTrue(hasattr(unpickled_instance, "validate")) + + +def clear_todo_controller_cache(): + """Helper method to clear controller cache for ToDo""" + if hasattr(frappe, "controllers") and frappe.local.site in frappe.controllers: + frappe.controllers[frappe.local.site].pop("ToDo", None)