diff --git a/CHANGELOG.md b/CHANGELOG.md index e7583d893..f54018b73 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 4b04786b9..f51fb3c6d 100644 --- a/CONTRIBUTORS.md +++ b/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) diff --git a/rich/console.py b/rich/console.py index 994adfc06..7800c2a58 100644 --- a/rich/console.py +++ b/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. diff --git a/tests/test_pickle_fix.py b/tests/test_pickle_fix.py new file mode 100644 index 000000000..e13fbea97 --- /dev/null +++ b/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()) diff --git a/tests/test_pickle_support.py b/tests/test_pickle_support.py new file mode 100644 index 000000000..c05f183d3 --- /dev/null +++ b/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)