Feature/pickle support

CHANGELOG.md

@@ -6,6 +6,12 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+## [Unreleased]
+
+### Added
+
+- Added pickle support to `ConsoleThreadLocals` and `Console` classes to enable serialization for caching frameworks https://github.com/Textualize/rich/pull/3853
+
 ## [14.1.0] - 2025-06-25
 
 ### Changed

CONTRIBUTORS.md

@@ -87,6 +87,7 @@ The following people have contributed to the development of Rich:
 - [James Addison](https://github.com/jayaddison)
 - [Pierro](https://github.com/xpierroz)
 - [Bernhard Wagner](https://github.com/bwagner)
+- [Tony Seah](https://github.com/pkusnail)
 - [Aaron Beaudoin](https://github.com/AaronBeaudoin)
 - [Sam Woodward](https://github.com/PyWoody)
 - [L. Yeung](https://github.com/lewis-yeung)

rich/console.py

@@ -546,6 +546,34 @@ class ConsoleThreadLocals(threading.local):
     buffer: List[Segment] = field(default_factory=list)
     buffer_index: int = 0
 
+    def __getstate__(self):
+        """Support for pickle serialization.
+
+        Returns the serializable state of the thread-local object.
+        Note: This loses the thread-local nature, but allows serialization
+        for caching and other use cases.
+
+        Returns:
+            Dict[str, Any]: The serializable state containing theme_stack,
+                           buffer, and buffer_index.
+        """
+        return {
+            "theme_stack": self.theme_stack,
+            "buffer": self.buffer.copy(),  # Create a copy to be safe
+            "buffer_index": self.buffer_index,
+        }
+
+    def __setstate__(self, state):
+        """Support for pickle deserialization.
+
+        Args:
+            state (Dict[str, Any]): The state dictionary from __getstate__
+        """
+        # Restore the state
+        self.theme_stack = state["theme_stack"]
+        self.buffer = state["buffer"]
+        self.buffer_index = state["buffer_index"]
+
 
 class RenderHook(ABC):
     """Provides hooks in to the render process."""
@@ -2611,6 +2639,39 @@ def save_svg(
         with open(path, "w", encoding="utf-8") as write_file:
             write_file.write(svg)
 
+    def __getstate__(self):
+        """Support for pickle serialization.
+
+        Returns the serializable state of the Console object.
+        Note: Thread locks are recreated during deserialization.
+
+        Returns:
+            Dict[str, Any]: The serializable state of the Console.
+        """
+        # Get all instance attributes except locks
+        state = self.__dict__.copy()
+
+        # Remove the unpickleable locks
+        state.pop("_lock", None)
+        state.pop("_record_buffer_lock", None)
+
+        return state
+
+    def __setstate__(self, state):
+        """Support for pickle deserialization.
+
+        Args:
+            state (Dict[str, Any]): The state dictionary from __getstate__
+        """
+        # Restore the state
+        self.__dict__.update(state)
+
+        # Recreate the locks
+        import threading
+
+        self._lock = threading.RLock()
+        self._record_buffer_lock = threading.RLock()
+
 
 def _svg_hash(svg_main_code: str) -> str:
     """Returns a unique hash for the given SVG main code.

tests/test_pickle_fix.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+"""
+Test script for ConsoleThreadLocals pickle support
+"""
+
+import pickle
+import sys
+import os
+
+# Add the current directory to the path so we can import the modified rich
+sys.path.insert(0, "/tmp/rich")
+
+from rich.console import Console
+from rich.segment import Segment
+
+
+def test_basic_pickle():
+    """Test basic pickle functionality of ConsoleThreadLocals."""
+    print("🧪 Testing basic ConsoleThreadLocals pickle functionality...")
+
+    console = Console()
+    ctl = console._thread_locals
+
+    # Add some data to make it more realistic
+    ctl.buffer.append(Segment("test"))
+    ctl.buffer_index = 1
+
+    try:
+        # Test serialization
+        pickled_data = pickle.dumps(ctl)
+        print("  ✅ Serialization successful")
+
+        # Test deserialization
+        restored_ctl = pickle.loads(pickled_data)
+        print("  ✅ Deserialization successful")
+
+        # Verify state preservation
+        assert type(restored_ctl.theme_stack) == type(ctl.theme_stack)
+        assert restored_ctl.buffer == ctl.buffer
+        assert restored_ctl.buffer_index == ctl.buffer_index
+        print("  ✅ State preservation verified")
+
+        return True
+
+    except Exception as e:
+        print(f"  ❌ Test failed: {e}")
+        return False
+
+
+def test_langflow_compatibility():
+    """Test compatibility with Langflow's caching mechanism."""
+    print("🔧 Testing Langflow cache compatibility...")
+
+    console = Console()
+
+    # Simulate Langflow's cache data structure
+    result_dict = {
+        "result": console,
+        "type": type(console),
+    }
+
+    try:
+        # This is what Langflow's cache service tries to do
+        pickled = pickle.dumps(result_dict)
+        print("  ✅ Complex object serialization successful")
+
+        restored = pickle.loads(pickled)
+        print("  ✅ Complex object deserialization successful")
+
+        # Verify the console is properly restored
+        assert type(restored["result"]) == type(console)
+        print("  ✅ Object type preservation verified")
+
+        return True
+
+    except Exception as e:
+        print(f"  ❌ Test failed: {e}")
+        return False
+
+
+def test_thread_local_behavior():
+    """Test that thread-local behavior works after unpickling."""
+    print("🔄 Testing thread-local behavior preservation...")
+
+    import threading
+    import time
+
+    console = Console()
+    ctl = console._thread_locals
+
+    # Serialize and deserialize
+    try:
+        pickled = pickle.dumps(ctl)
+        restored_ctl = pickle.loads(pickled)
+
+        # Test that we can still use the restored object
+        restored_ctl.buffer.append(Segment("thread test"))
+        restored_ctl.buffer_index = 5
+
+        print(f"  ✅ Restored object is functional")
+        print(f"  📊 Buffer length: {len(restored_ctl.buffer)}")
+        print(f"  📊 Buffer index: {restored_ctl.buffer_index}")
+
+        return True
+
+    except Exception as e:
+        print(f"  ❌ Test failed: {e}")
+        return False
+
+
+def main():
+    """Run all tests."""
+    print("🚀 Starting Rich ConsoleThreadLocals pickle fix tests...\n")
+
+    tests = [
+        test_basic_pickle,
+        test_langflow_compatibility,
+        test_thread_local_behavior,
+    ]
+
+    passed = 0
+    total = len(tests)
+
+    for test in tests:
+        if test():
+            passed += 1
+        print()  # Add spacing between tests
+
+    print("=" * 50)
+    print(f"📊 Test Results: {passed}/{total} tests passed")
+
+    if passed == total:
+        print("🎉 All tests passed! The pickle fix is working correctly.")
+        return 0
+    else:
+        print("❌ Some tests failed. Please check the output above.")
+        return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

tests/test_pickle_support.py

@@ -0,0 +1,155 @@
+"""Tests for pickle support in Rich objects."""
+
+import pickle
+from rich.console import Console, ConsoleThreadLocals
+from rich.segment import Segment
+from rich.theme import Theme, ThemeStack
+
+
+def test_console_thread_locals_pickle():
+    """Test that ConsoleThreadLocals can be pickled and unpickled."""
+    console = Console()
+    ctl = console._thread_locals
+
+    # Add some data to make it more realistic
+    ctl.buffer.append(Segment("test"))
+    ctl.buffer_index = 1
+
+    # Test serialization
+    pickled_data = pickle.dumps(ctl)
+
+    # Test deserialization
+    restored_ctl = pickle.loads(pickled_data)
+
+    # Verify state preservation
+    assert type(restored_ctl.theme_stack) == type(ctl.theme_stack)
+    assert restored_ctl.buffer == ctl.buffer
+    assert restored_ctl.buffer_index == ctl.buffer_index
+
+
+def test_console_pickle():
+    """Test that Console objects can be pickled and unpickled."""
+    console = Console(width=120, height=40)
+
+    # Test serialization
+    pickled_data = pickle.dumps(console)
+
+    # Test deserialization
+    restored_console = pickle.loads(pickled_data)
+
+    # Verify basic properties are preserved
+    assert restored_console.width == console.width
+    assert restored_console.height == console.height
+    assert restored_console._color_system == console._color_system
+
+    # Verify locks are recreated
+    assert hasattr(restored_console, "_lock")
+    assert hasattr(restored_console, "_record_buffer_lock")
+
+    # Verify the console is functional
+    with restored_console.capture() as capture:
+        restored_console.print("Test message")
+
+    assert "Test message" in capture.get()
+
+
+def test_console_with_complex_state_pickle():
+    """Test console pickle with more complex state."""
+    theme = Theme({"info": "cyan", "warning": "yellow", "error": "red bold"})
+
+    console = Console(theme=theme, record=True)
+
+    # Add some content
+    console.print("Info message", style="info")
+    console.print("Warning message", style="warning")
+    console.record = False  # Stop recording
+
+    # Test serialization
+    pickled_data = pickle.dumps(console)
+
+    # Test deserialization
+    restored_console = pickle.loads(pickled_data)
+
+    # Verify theme is preserved
+    assert restored_console.get_style("info").color.name == "cyan"
+    assert restored_console.get_style("warning").color.name == "yellow"
+
+    # Verify console functionality
+    assert restored_console.record is False
+
+
+def test_cache_simulation():
+    """Test cache-like usage scenario (similar to Langflow)."""
+    console = Console()
+
+    # Simulate caching scenario like Langflow
+    cache_data = {
+        "result": console,
+        "type": type(console),
+        "metadata": {"created": "2025-09-25", "version": "1.0"},
+    }
+
+    # This should not raise any pickle errors
+    pickled = pickle.dumps(cache_data)
+    restored = pickle.loads(pickled)
+
+    # Verify restoration
+    assert type(restored["result"]) == Console
+    assert restored["type"] == Console
+    assert restored["metadata"]["created"] == "2025-09-25"
+
+    # Verify the restored console works
+    restored_console = restored["result"]
+    with restored_console.capture() as capture:
+        restored_console.print("Cache test successful")
+
+    assert "Cache test successful" in capture.get()
+
+
+def test_nested_console_pickle():
+    """Test pickling dict containing Console instances."""
+    # Use a simple dict instead of local class to avoid pickle issues
+    container = {
+        "console": Console(width=100),
+        "name": "test_container",
+        "data": [1, 2, 3],
+    }
+
+    # Should be able to pickle dict containing Console
+    pickled = pickle.dumps(container)
+    restored = pickle.loads(pickled)
+
+    assert restored["name"] == "test_container"
+    assert restored["data"] == [1, 2, 3]
+    assert restored["console"].width == 100
+
+    # Verify console functionality
+    with restored["console"].capture() as capture:
+        restored["console"].print("Nested test")
+
+    assert "Nested test" in capture.get()
+
+
+if __name__ == "__main__":
+    # Run tests manually if called directly
+    import sys
+
+    tests = [
+        test_console_thread_locals_pickle,
+        test_console_pickle,
+        test_console_with_complex_state_pickle,
+        test_cache_simulation,
+        test_nested_console_pickle,
+    ]
+
+    passed = 0
+    for test in tests:
+        try:
+            test()
+            print(f"✅ {test.__name__} passed")
+            passed += 1
+        except Exception as e:
+            print(f"❌ {test.__name__} failed: {e}")
+
+    print(f"\n📊 Results: {passed}/{len(tests)} tests passed")
+    sys.exit(0 if passed == len(tests) else 1)