From 86b742e91dcd4b86c319800d0422a76ad6a3d818 Mon Sep 17 00:00:00 2001
From: Kenton Varda <temporal@gmail.com>
Date: Sat, 31 Aug 2013 17:07:44 -0700
Subject: [PATCH] Wrap MessageBuilder::setRoot().

---
 capnp/capnp.pyx         | 17 +++++++++++++++++
 capnp/schema_cpp.pxd    |  1 +
 test/test_regression.py |  9 +++++++++
 3 files changed, 27 insertions(+)

diff --git a/capnp/capnp.pyx b/capnp/capnp.pyx
index 9c19c73..827a213 100644
--- a/capnp/capnp.pyx
+++ b/capnp/capnp.pyx
@@ -32,6 +32,10 @@ ctypedef double Float64
 from libc.stdlib cimport malloc, free
 from libcpp cimport bool as cbool
 
+ctypedef fused _DynamicStructReaderOrBuilder:
+    _DynamicStructReader
+    _DynamicStructBuilder
+
 def _make_enum(enum_name, *sequential, **named):
     enums = dict(zip(sequential, range(len(sequential))), **named)
     reverse = dict((value, key) for key, value in enums.iteritems())
@@ -971,6 +975,19 @@ cdef class MessageBuilder:
         else:
             s = schema
         return _DynamicStructBuilder()._init(self.thisptr.getRootDynamicStruct(s.thisptr), self)
+    
+    cpdef setRoot(self, value):
+        """A method for instantiating Cap'n Proto structs by copying from an existing struct
+
+        :type value: :class:`_DynamicStructReader`
+        :param value: A Cap'n Proto struct value to copy
+
+        :rtype: void
+        """
+        
+        if type(value) is _DynamicStructBuilder:
+            value = value.asReader();
+        self.thisptr.setRootDynamicStruct((<_DynamicStructReader>value).thisptr)
 
     cpdef newOrphan(self, schema):
         """A method for instantiating Cap'n Proto orphans
diff --git a/capnp/schema_cpp.pxd b/capnp/schema_cpp.pxd
index f9a56b8..c8dcc7b 100644
--- a/capnp/schema_cpp.pxd
+++ b/capnp/schema_cpp.pxd
@@ -664,6 +664,7 @@ cdef extern from "capnp/message.h" namespace " ::capnp":
         
         DynamicStruct.Builder getRootDynamicStruct'getRoot< ::capnp::DynamicStruct>'(StructSchema)
         DynamicStruct.Builder initRootDynamicStruct'initRoot< ::capnp::DynamicStruct>'(StructSchema)
+        void setRootDynamicStruct'setRoot< ::capnp::DynamicStruct::Reader>'(DynamicStruct.Reader)
 
         DynamicOrphan newOrphan'getOrphanage().newOrphan'(StructSchema)
 
diff --git a/test/test_regression.py b/test/test_regression.py
index 8db871b..f19683e 100644
--- a/test/test_regression.py
+++ b/test/test_regression.py
@@ -281,6 +281,15 @@ def test_binary_read(all_types):
     expectedText = open(os.path.join(this_dir, 'all-types.txt'), 'r').read()
     assert str(root) + '\n' == expectedText
 
+    # Test setRoot().
+    builder = capnp.MallocMessageBuilder()
+    builder.setRoot(root)
+    check_all_types(builder.getRoot(all_types.TestAllTypes))
+
+    builder2 = capnp.MallocMessageBuilder()
+    builder2.setRoot(builder.getRoot(all_types.TestAllTypes))
+    check_all_types(builder2.getRoot(all_types.TestAllTypes))
+
 def test_packed_read(all_types):
     f = open(os.path.join(this_dir, 'all-types.packed'), 'r')
     message = capnp.PackedFdMessageReader(f.fileno())