feat: support pickling extended class

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
 

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)