feat(recipe): add support for curator SharedCount recipe

kazoo/recipe/counter.py

@@ -6,7 +6,7 @@
 """
 from kazoo.exceptions import BadVersionError
 from kazoo.retry import ForceRetryError
-
+import struct
 
 class Counter(object):
     """Kazoo Counter
@@ -19,6 +19,12 @@ class Counter(object):
     `type(counter.default)(value)` both using an ascii encoding. As
     such other data types might be used for the counter value.
 
+    If you would like to support clients updating the same znode path using
+    either kazoo's counter recipe or curator's SharedCount recipe, you will
+    need to enable the support_curator flag. This flag limits
+    support to integers only and does not use ascii encoding as described
+    above.
+
     Counter changes can raise
     :class:`~kazoo.exceptions.BadVersionError` if the retry policy
     wasn't able to apply a change.
@@ -42,22 +48,35 @@ class Counter(object):
         counter.pre_value == 1.0
         counter.post_value == 3.0
 
+        counter = zk.Counter("/curator", support_curator=True)
+        counter += 2
+        counter -= 1
+        counter.value == 1
+        counter.pre_value == 2
+        counter.post_value == 1
+
     """
-    def __init__(self, client, path, default=0):
+    def __init__(self, client, path, default=0, support_curator=False):
         """Create a Kazoo Counter
 
         :param client: A :class:`~kazoo.client.KazooClient` instance.
         :param path: The counter path to use.
-        :param default: The default value.
+        :param default: The default value to use for new counter paths.
+        :param support_curator: Enable if support for curator's SharedCount
+                                recipe is desired.
 
         """
         self.client = client
         self.path = path
         self.default = default
         self.default_type = type(default)
+        self.support_curator = support_curator
         self._ensured_path = False
         self.pre_value = None
         self.post_value = None
+        if self.support_curator and not isinstance(self.default, int):
+            raise TypeError("when support_curator is enabled the default "
+                            "type must be an int")
 
     def _ensure_node(self):
         if not self._ensured_path:
@@ -68,7 +87,10 @@ def _ensure_node(self):
     def _value(self):
         self._ensure_node()
         old, stat = self.client.get(self.path)
-        old = old.decode('ascii') if old != b'' else self.default
+        if self.support_curator:
+            old = struct.unpack(">i", old)[0] if old != b'' else self.default
+        else:
+            old = old.decode('ascii') if old != b'' else self.default
         version = stat.version
         data = self.default_type(old)
         return data, version
@@ -86,7 +108,10 @@ def _change(self, value):
     def _inner_change(self, value):
         self.pre_value, version = self._value()
         post_value = self.pre_value + value
-        data = repr(post_value).encode('ascii')
+        if self.support_curator:
+            data = struct.pack(">i", post_value)
+        else:
+            data = repr(post_value).encode('ascii')
         try:
             self.client.set(self.path, data, version=version)
         except BadVersionError:  # pragma: nocover

kazoo/tests/test_counter.py

@@ -21,6 +21,22 @@ def test_int_counter(self):
         counter - 1
         eq_(counter.value, -1)
 
+    def test_int_curator_counter(self):
+        counter = self._makeOne(support_curator=True)
+        eq_(counter.value, 0)
+        counter += 2
+        counter + 1
+        eq_(counter.value, 3)
+        counter -= 3
+        counter - 1
+        eq_(counter.value, -1)
+        counter += 1
+        counter += 2147483647
+        eq_(counter.value, 2147483647)
+        counter -= 2147483647
+        counter -= 2147483647
+        eq_(counter.value, -2147483647)
+
     def test_float_counter(self):
         counter = self._makeOne(default=0.0)
         eq_(counter.value, 0.0)
@@ -33,6 +49,8 @@ def test_errors(self):
         counter = self._makeOne()
         self.assertRaises(TypeError, counter.__add__, 2.1)
         self.assertRaises(TypeError, counter.__add__, b"a")
+        with self.assertRaises(TypeError):
+            counter = self._makeOne(default=0.0, support_curator=True)
 
     def test_pre_post_values(self):
         counter = self._makeOne()