Make histograms pickleable

hdrh/histogram.py

@@ -137,6 +137,56 @@ def __init__(self,
         # no tag by default
         self.tag = None
 
+    def __getstate__(self):
+        '''Return a pickle-safe representation of this histogram.
+
+        The live encoder owns ctypes-backed count storage, so serializing the
+        default instance dictionary is not portable across pickle round trips.
+        '''
+        counts = []
+        for index in range(self.counts_len):
+            count = self.get_count_at_index(index)
+            if count:
+                counts.append((index, count))
+
+        return {
+            'lowest_trackable_value': self.lowest_trackable_value,
+            'highest_trackable_value': self.highest_trackable_value,
+            'significant_figures': self.significant_figures,
+            'word_size': self.word_size,
+            'b64_wrap': self.b64_wrap,
+            'int_to_double_conversion_ratio': self.int_to_double_conversion_ratio,
+            'min_value': self.min_value,
+            'max_value': self.max_value,
+            'total_count': self.total_count,
+            'start_time_stamp_msec': self.start_time_stamp_msec,
+            'end_time_stamp_msec': self.end_time_stamp_msec,
+            'tag': self.tag,
+            'counts': counts,
+        }
+
+    def __setstate__(self, state):
+        '''Restore a histogram serialized by __getstate__.'''
+        self.__init__(state['lowest_trackable_value'],
+                      state['highest_trackable_value'],
+                      state['significant_figures'],
+                      word_size=state.get('word_size', 8),
+                      b64_wrap=state.get('b64_wrap', True))
+
+        for index, count in state.get('counts', ()):
+            self.counts[index] = count
+
+        self.int_to_double_conversion_ratio = \
+            state.get('int_to_double_conversion_ratio', 1.0)
+        self.encoder.payload.payload.conversion_ratio_bits = \
+            self.int_to_double_conversion_ratio
+        self.min_value = state.get('min_value', sys.maxsize)
+        self.max_value = state.get('max_value', 0)
+        self.total_count = state.get('total_count', 0)
+        self.start_time_stamp_msec = state.get('start_time_stamp_msec', 0)
+        self.end_time_stamp_msec = state.get('end_time_stamp_msec', 0)
+        self.tag = state.get('tag')
+
     def _clz(self, value):
         """calculate the leading zeros, equivalent to C __builtin_clzll()
         value in hex:

test/test_hdrhistogram.py

@@ -25,6 +25,7 @@
 import datetime
 import io
 import os
+import pickle
 import struct
 import zlib
 import sys
@@ -85,6 +86,26 @@ def test_empty_histogram():
     assert histogram.get_mean_value() == 0
     assert histogram.get_stddev() == 0
 
+@pytest.mark.basic
+def test_pickle_round_trip():
+    histogram = HdrHistogram(LOWEST, HIGHEST, SIGNIFICANT, word_size=4,
+                             b64_wrap=False)
+    histogram.record_value(1, 3)
+    histogram.record_value(2000, 4)
+    histogram.set_start_time_stamp(123)
+    histogram.set_end_time_stamp(456)
+    histogram.set_tag('pickle')
+
+    restored = pickle.loads(pickle.dumps(histogram))
+
+    assert restored is not histogram
+    assert restored.equals(histogram)
+    assert restored.get_word_size() == histogram.get_word_size()
+    assert restored.b64_wrap == histogram.b64_wrap
+    assert restored.get_start_time_stamp() == 123
+    assert restored.get_end_time_stamp() == 456
+    assert restored.get_tag() == 'pickle'
+
 @pytest.mark.basic
 def test_large_numbers():
     histogram = HdrHistogram(20000000, 100000000, 5)