From 0e395b3e728a43febcf1e1cc0873d4a75ef20676 Mon Sep 17 00:00:00 2001 From: pkusnail Date: Thu, 25 Sep 2025 16:02:51 +0800 Subject: [PATCH 1/6] feat: Add pickle support to ConsoleThreadLocals and Console classes - Add __getstate__ and __setstate__ methods to ConsoleThreadLocals class - Add __getstate__ and __setstate__ methods to Console class - Enable serialization for caching frameworks like Redis - Maintains backward compatibility and thread safety - Fixes 'cannot pickle ConsoleThreadLocals object' error --- rich/console.py | 60 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 60 insertions(+) diff --git a/rich/console.py b/rich/console.py index 994adfc069..8c462d543a 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.""" @@ -2610,6 +2638,38 @@ 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: From 0428a42be37418cea2ba76a73332ff4352d9ab95 Mon Sep 17 00:00:00 2001 From: pkusnail Date: Thu, 25 Sep 2025 16:04:12 +0800 Subject: [PATCH 2/6] test: Add comprehensive pickle support tests - Add test_pickle_support.py with 5 comprehensive test cases - Test ConsoleThreadLocals and Console pickle functionality - Test cache simulation scenario (Langflow compatibility) - Test complex state preservation and nested objects - All tests pass and verify proper serialization/deserialization --- tests/test_pickle_support.py | 159 +++++++++++++++++++++++++++++++++++ 1 file changed, 159 insertions(+) create mode 100644 tests/test_pickle_support.py diff --git a/tests/test_pickle_support.py b/tests/test_pickle_support.py new file mode 100644 index 0000000000..660b08f541 --- /dev/null +++ b/tests/test_pickle_support.py @@ -0,0 +1,159 @@ +"""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) \ No newline at end of file From ccdfeff27040aaf93c115e8dc55207a030a83c65 Mon Sep 17 00:00:00 2001 From: pkusnail Date: Thu, 25 Sep 2025 16:04:45 +0800 Subject: [PATCH 3/6] test: Add initial pickle fix validation script - Add test_pickle_fix.py for manual testing and validation - Tests basic functionality, Langflow compatibility, and thread-local behavior - Provides detailed output for debugging and verification - Can be run independently for quick validation --- tests/test_pickle_fix.py | 141 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 141 insertions(+) create mode 100644 tests/test_pickle_fix.py diff --git a/tests/test_pickle_fix.py b/tests/test_pickle_fix.py new file mode 100644 index 0000000000..be3365e986 --- /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()) \ No newline at end of file From d97d8b26c7d8c01dd7c2c3b32ea38b2eab5bea26 Mon Sep 17 00:00:00 2001 From: pkusnail Date: Thu, 25 Sep 2025 16:19:36 +0800 Subject: [PATCH 4/6] docs: Add pickle support entry to CHANGELOG.md - Add entry under [Unreleased] section for pickle support feature - Reference PR #3853 for ConsoleThreadLocals and Console serialization - Follows Keep a Changelog format --- CHANGELOG.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index e7583d8936..f54018b737 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 From 7f32d7f52794e8ec8b5724ae2ab11d73c235fa83 Mon Sep 17 00:00:00 2001 From: pkusnail Date: Thu, 25 Sep 2025 16:20:51 +0800 Subject: [PATCH 5/6] docs: Add Tony Seah to CONTRIBUTORS.md - Add Tony Seah (pkusnail) to contributors list for pickle support contribution - Alphabetically placed for proper organization --- CONTRIBUTORS.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CONTRIBUTORS.md b/CONTRIBUTORS.md index 4b04786b9c..f51fb3c6d5 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) From 3f2eb2d988fe22e3598542dd1773ae010ea4aacd Mon Sep 17 00:00:00 2001 From: pkusnail Date: Thu, 25 Sep 2025 16:23:45 +0800 Subject: [PATCH 6/6] style: Apply black formatting to pickle support code - Format rich/console.py with black for consistency - Format tests/test_pickle_support.py and test_pickle_fix.py - Ensures code style compliance with Rich project standards --- rich/console.py | 41 ++++++++++---------- tests/test_pickle_fix.py | 56 ++++++++++++++-------------- tests/test_pickle_support.py | 72 +++++++++++++++++------------------- 3 files changed, 83 insertions(+), 86 deletions(-) diff --git a/rich/console.py b/rich/console.py index 8c462d543a..7800c2a582 100644 --- a/rich/console.py +++ b/rich/console.py @@ -548,31 +548,31 @@ class ConsoleThreadLocals(threading.local): 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 + "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'] + self.theme_stack = state["theme_stack"] + self.buffer = state["buffer"] + self.buffer_index = state["buffer_index"] class RenderHook(ABC): @@ -2638,36 +2638,37 @@ 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) - + 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() diff --git a/tests/test_pickle_fix.py b/tests/test_pickle_fix.py index be3365e986..e13fbea97f 100644 --- a/tests/test_pickle_fix.py +++ b/tests/test_pickle_fix.py @@ -8,7 +8,7 @@ import os # Add the current directory to the path so we can import the modified rich -sys.path.insert(0, '/tmp/rich') +sys.path.insert(0, "/tmp/rich") from rich.console import Console from rich.segment import Segment @@ -17,31 +17,31 @@ 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 @@ -50,29 +50,29 @@ def test_basic_pickle(): 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 @@ -81,28 +81,28 @@ def test_langflow_compatibility(): 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 @@ -111,24 +111,24 @@ def test_thread_local_behavior(): 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 @@ -138,4 +138,4 @@ def main(): if __name__ == "__main__": - sys.exit(main()) \ No newline at end of file + sys.exit(main()) diff --git a/tests/test_pickle_support.py b/tests/test_pickle_support.py index 660b08f541..c05f183d3f 100644 --- a/tests/test_pickle_support.py +++ b/tests/test_pickle_support.py @@ -10,17 +10,17 @@ 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 @@ -30,54 +30,50 @@ def test_console_thread_locals_pickle(): 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') - + 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" - }) - + 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 @@ -85,28 +81,28 @@ def test_console_with_complex_state_pickle(): 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"} + "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() @@ -116,28 +112,28 @@ def test_nested_console_pickle(): container = { "console": Console(width=100), "name": "test_container", - "data": [1, 2, 3] + "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, @@ -145,7 +141,7 @@ def test_nested_console_pickle(): test_cache_simulation, test_nested_console_pickle, ] - + passed = 0 for test in tests: try: @@ -154,6 +150,6 @@ def test_nested_console_pickle(): 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) \ No newline at end of file + sys.exit(0 if passed == len(tests) else 1)