feat(core): Add support for Container and TTL nodes

Also add support through transations.

Closes #334, #496

Author: ceache

.flake8

@@ -1,2 +1,2 @@
 [flake8]
-ignore = BLK100
+ignore = BLK100,W503

kazoo/client.py

@@ -32,6 +32,8 @@
     CloseInstance,
     Create,
     Create2,
+    CreateContainer,
+    CreateTTL,
     Delete,
     Exists,
     GetChildren,
@@ -873,7 +875,8 @@ def sync(self, path):
         return self.sync_async(path).get()
 
     def create(self, path, value=b"", acl=None, ephemeral=False,
-               sequence=False, makepath=False, include_data=False):
+               sequence=False, makepath=False, include_data=False,
+               container=False, ttl=0):
         """Create a node with the given value as its data. Optionally
         set an ACL on the node.
 
@@ -950,15 +953,19 @@ def create(self, path, value=b"", acl=None, ephemeral=False,
             The `makepath` option.
         .. versionadded:: 2.7
             The `include_data` option.
+        .. versionadded:: 2.8
+            The container and ttl options.
+
         """
         acl = acl or self.default_acl
         return self.create_async(
             path, value, acl=acl, ephemeral=ephemeral,
-            sequence=sequence, makepath=makepath, include_data=include_data
-        ).get()
+            sequence=sequence, makepath=makepath, include_data=include_data,
+            container=container, ttl=ttl).get()
 
     def create_async(self, path, value=b"", acl=None, ephemeral=False,
-                     sequence=False, makepath=False, include_data=False):
+                     sequence=False, makepath=False, include_data=False,
+                     container=False, ttl=0):
         """Asynchronously create a ZNode. Takes the same arguments as
         :meth:`create`.
 
@@ -967,45 +974,35 @@ def create_async(self, path, value=b"", acl=None, ephemeral=False,
         .. versionadded:: 1.1
             The makepath option.
         .. versionadded:: 2.7
-            The `include_data` option.
+            The include_data option.
+        .. versionadded:: 2.8
+            The container and ttl options.
         """
         if acl is None and self.default_acl:
             acl = self.default_acl
-
-        if not isinstance(path, string_types):
-            raise TypeError("Invalid type for 'path' (string expected)")
-        if acl and (isinstance(acl, ACL) or
-                    not isinstance(acl, (tuple, list))):
-            raise TypeError("Invalid type for 'acl' (acl must be a tuple/list"
-                            " of ACL's")
-        if value is not None and not isinstance(value, bytes_types):
-            raise TypeError("Invalid type for 'value' (must be a byte string)")
-        if not isinstance(ephemeral, bool):
-            raise TypeError("Invalid type for 'ephemeral' (bool expected)")
-        if not isinstance(sequence, bool):
-            raise TypeError("Invalid type for 'sequence' (bool expected)")
         if not isinstance(makepath, bool):
             raise TypeError("Invalid type for 'makepath' (bool expected)")
-        if not isinstance(include_data, bool):
-            raise TypeError("Invalid type for 'include_data' (bool expected)")
-
-        flags = 0
-        if ephemeral:
-            flags |= 1
-        if sequence:
-            flags |= 2
-        if acl is None:
-            acl = OPEN_ACL_UNSAFE
 
+        opcode = _create_opcode(
+            path, value, acl, self.chroot,
+            ephemeral, sequence, include_data, container, ttl
+        )
         async_result = self.handler.async_result()
 
         @capture_exceptions(async_result)
         def do_create():
-            result = self._create_async_inner(
-                path, value, acl, flags,
-                trailing=sequence, include_data=include_data
-            )
-            result.rawlink(create_completion)
+            inner_async_result = self.handler.async_result()
+
+            call_result = self._call(opcode, inner_async_result)
+            if call_result is False:
+                # We hit a short-circuit exit on the _call. Because we are
+                # not using the original async_result here, we bubble the
+                # exception upwards to the do_create function in
+                # KazooClient.create so that it gets set on the correct
+                # async_result object
+                raise inner_async_result.exception
+
+            inner_async_result.rawlink(create_completion)
 
         @capture_exceptions(async_result)
         def retry_completion(result):
@@ -1015,11 +1012,11 @@ def retry_completion(result):
         @wrap(async_result)
         def create_completion(result):
             try:
-                if include_data:
+                if opcode.type == Create.type:
+                    return self.unchroot(result.get())
+                else:
                     new_path, stat = result.get()
                     return self.unchroot(new_path), stat
-                else:
-                    return self.unchroot(result.get())
             except NoNodeError:
                 if not makepath:
                     raise
@@ -1032,26 +1029,6 @@ def create_completion(result):
         do_create()
         return async_result
 
-    def _create_async_inner(self, path, value, acl, flags,
-                            trailing=False, include_data=False):
-        async_result = self.handler.async_result()
-        if include_data:
-            opcode = Create2
-        else:
-            opcode = Create
-
-        call_result = self._call(
-            opcode(_prefix_root(self.chroot, path, trailing=trailing),
-                   value, acl, flags), async_result)
-        if call_result is False:
-            # We hit a short-circuit exit on the _call. Because we are
-            # not using the original async_result here, we bubble the
-            # exception upwards to the do_create function in
-            # KazooClient.create so that it gets set on the correct
-            # async_result object
-            raise async_result.exception
-        return async_result
-
     def ensure_path(self, path, acl=None):
         """Recursively create a path if it doesn't exist.
 
@@ -1590,39 +1567,24 @@ def __init__(self, client):
         self.committed = False
 
     def create(self, path, value=b"", acl=None, ephemeral=False,
-               sequence=False):
+               sequence=False, include_data=False, container=False, ttl=0):
         """Add a create ZNode to the transaction. Takes the same
         arguments as :meth:`KazooClient.create`, with the exception
         of `makepath`.
 
         :returns: None
 
+        .. versionadded:: 2.8
+            The include_data, container and ttl options.
         """
         if acl is None and self.client.default_acl:
             acl = self.client.default_acl
 
-        if not isinstance(path, string_types):
-            raise TypeError("Invalid type for 'path' (string expected)")
-        if acl and not isinstance(acl, (tuple, list)):
-            raise TypeError("Invalid type for 'acl' (acl must be a tuple/list"
-                            " of ACL's")
-        if not isinstance(value, bytes_types):
-            raise TypeError("Invalid type for 'value' (must be a byte string)")
-        if not isinstance(ephemeral, bool):
-            raise TypeError("Invalid type for 'ephemeral' (bool expected)")
-        if not isinstance(sequence, bool):
-            raise TypeError("Invalid type for 'sequence' (bool expected)")
-
-        flags = 0
-        if ephemeral:
-            flags |= 1
-        if sequence:
-            flags |= 2
-        if acl is None:
-            acl = OPEN_ACL_UNSAFE
-
-        self._add(Create(_prefix_root(self.client.chroot, path), value, acl,
-                         flags), None)
+        opcode = _create_opcode(
+            path, value, acl, self.client.chroot,
+            ephemeral, sequence, include_data, container, ttl
+        )
+        self._add(opcode, None)
 
     def delete(self, path, version=-1):
         """Add a delete ZNode to the transaction. Takes the same
@@ -1701,3 +1663,85 @@ def _add(self, request, post_processor=None):
         self._check_tx_state()
         self.client.logger.log(BLATHER, 'Added %r to %r', request, self)
         self.operations.append(request)
+
+
+def _create_opcode(path, value, acl, chroot,
+                   ephemeral, sequence, include_data, container, ttl):
+        """Helper function.
+        Creates the create OpCode for regular `client.create()` operations as
+        well as in a `client.transaction()` context.
+        """
+        if not isinstance(path, string_types):
+            raise TypeError("Invalid type for 'path' (string expected)")
+        if acl and (isinstance(acl, ACL) or
+                    not isinstance(acl, (tuple, list))):
+            raise TypeError("Invalid type for 'acl' (acl must be a tuple/list"
+                            " of ACL's")
+        if value is not None and not isinstance(value, bytes_types):
+            raise TypeError("Invalid type for 'value' (must be a byte string)")
+        if not isinstance(ephemeral, bool):
+            raise TypeError("Invalid type for 'ephemeral' (bool expected)")
+        if not isinstance(sequence, bool):
+            raise TypeError("Invalid type for 'sequence' (bool expected)")
+        if not isinstance(include_data, bool):
+            raise TypeError("Invalid type for 'include_data' (bool expected)")
+        if not isinstance(container, bool):
+            raise TypeError("Invalid type for 'container' (bool expected)")
+        if not isinstance(ttl, int) or ttl < 0:
+            raise TypeError("Invalid 'ttl' (integer >= 0 expected)")
+        if ttl and ephemeral:
+            raise TypeError("Invalid node creation: ephemeral & ttl")
+        if container and (ephemeral or sequence or ttl):
+            raise TypeError(
+                "Invalid node creation: container & ephemeral/sequence/ttl"
+            )
+
+        # Should match Zookeeper's CreateMode fromFlag
+        # https://github.com/apache/zookeeper/blob/master/zookeeper-server/
+        # src/main/java/org/apache/zookeeper/CreateMode.java#L112
+        flags = 0
+        if ephemeral:
+            flags |= 1
+        if sequence:
+            flags |= 2
+        if container:
+            flags = 4
+        if ttl:
+            if sequence:
+                flags = 6
+            else:
+                flags = 5
+
+        if acl is None:
+            acl = OPEN_ACL_UNSAFE
+
+        # Figure out the OpCode we are going to send
+        if include_data:
+            return Create2(
+                _prefix_root(chroot, path, trailing=sequence),
+                value,
+                acl,
+                flags
+            )
+        elif container:
+            return CreateContainer(
+                _prefix_root(chroot, path, trailing=False),
+                value,
+                acl,
+                flags
+            )
+        elif ttl:
+            return CreateTTL(
+                _prefix_root(chroot, path, trailing=sequence),
+                value,
+                acl,
+                flags,
+                ttl
+            )
+        else:
+            return Create(
+                _prefix_root(chroot, path, trailing=sequence),
+                value,
+                acl,
+                flags
+            )