diff --git a/hdrh/histogram.py b/hdrh/histogram.py index d145327..59832a6 100644 --- a/hdrh/histogram.py +++ b/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: diff --git a/test/test_hdrhistogram.py b/test/test_hdrhistogram.py index 01ceb2f..7922a6a 100644 --- a/test/test_hdrhistogram.py +++ b/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)