Added __getstate__ and __setstate__ for pickling

svpcom · Feb 20, 2014 · 9d74d93 · 9d74d93
1 parent f91ed5d
commit 9d74d93
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 0 deletions.
diff --git a/hyperloglog/hll.py b/hyperloglog/hll.py
@@ -76,6 +76,13 @@ def __init__(self, error_rate):
         self.m = 1 << p
         self.M = [ 0 for i in range(self.m) ]
 
+    def __getstate__(self):
+        return dict([x, getattr(self, x)] for x in self.__slots__)
+
+    def __setstate__(self, d):
+        for key in d:
+            setattr(self, key, d[key])
+
     def add(self, value):
         """
         Adds the item to the HyperLogLog

diff --git a/hyperloglog/shll.py b/hyperloglog/shll.py
@@ -47,6 +47,13 @@ def __init__(self, error_rate, window, lpfm=None):
         self.p = p
         self.m = m
 
+    def __getstate__(self):
+        return dict([x, getattr(self, x)] for x in self.__slots__)
+
+    def __setstate__(self, d):
+        for key in d:
+            setattr(self, key, d[key])
+
     @classmethod
     def from_list(cls, lpfm, window):
         return cls(None, window, lpfm)