Merge pull request #363 from danring/fix-341

Emit GetRootAs methods for all types in Go and Python

9 files changed, 80 insertions, 10 deletions

diff --git a/src/idl_gen_go.cpp b/src/idl_gen_go.cpp
index 0920b772..1f6b103e 100644
--- a/src/idl_gen_go.cpp
+++ b/src/idl_gen_go.cpp
@@ -551,13 +551,12 @@ static void GenTableBuilders(const StructDef &struct_def,
 
 // Generate struct or table methods.
 static void GenStruct(const StructDef &struct_def,
-                      std::string *code_ptr,
-                      StructDef *root_struct_def) {
+                      std::string *code_ptr) {
   if (struct_def.generated) return;
 
   GenComment(struct_def.doc_comment, code_ptr, nullptr);
   BeginClass(struct_def, code_ptr);
-  if (&struct_def == root_struct_def) {
+  if (!struct_def.fixed) {
     // Generate a special accessor for the table that has been declared as
     // the root type.
     NewRootTypeFromBuffer(struct_def, code_ptr);
@@ -681,7 +680,7 @@ class GoGenerator : public BaseGenerator {
     for (auto it = parser_.structs_.vec.begin();
          it != parser_.structs_.vec.end(); ++it) {
       std::string declcode;
-      go::GenStruct(**it, &declcode, parser_.root_struct_def_);
+      go::GenStruct(**it, &declcode);
       if (!SaveType(**it, declcode, true)) return false;
     }
 
diff --git a/src/idl_gen_python.cpp b/src/idl_gen_python.cpp
index 7e8dc0eb..52944bfc 100644
--- a/src/idl_gen_python.cpp
+++ b/src/idl_gen_python.cpp
@@ -94,7 +94,7 @@ static void NewRootTypeFromBuffer(const StructDef &struct_def,
   code += Indent + Indent + "x = " + struct_def.name + "()\n";
   code += Indent + Indent + "x.Init(buf, n + offset)\n";
   code += Indent + Indent + "return x\n";
-  code += "\n\n";
+  code += "\n";
 }
 
 // Initialize an existing object with other data, to avoid an allocation.
@@ -478,13 +478,12 @@ static void GenTableBuilders(const StructDef &struct_def,
 
 // Generate struct or table methods.
 static void GenStruct(const StructDef &struct_def,
-                      std::string *code_ptr,
-                      StructDef *root_struct_def) {
+                      std::string *code_ptr) {
   if (struct_def.generated) return;
 
   GenComment(struct_def.doc_comment, code_ptr, nullptr, "# ");
   BeginClass(struct_def, code_ptr);
-  if (&struct_def == root_struct_def) {
+  if (!struct_def.fixed) {
     // Generate a special accessor for the table that has been declared as
     // the root type.
     NewRootTypeFromBuffer(struct_def, code_ptr);
@@ -621,7 +620,7 @@ class PythonGenerator : public BaseGenerator {
          it != parser_.structs_.vec.end(); ++it) {
       auto &struct_def = **it;
       std::string declcode;
-      GenStruct(struct_def, &declcode, parser_.root_struct_def_);
+      GenStruct(struct_def, &declcode);
       if (!SaveType(struct_def, declcode, true)) return false;
     }
     return true;
diff --git a/tests/MyGame/Example/Monster.py b/tests/MyGame/Example/Monster.py
index 9d2909e8..a7a78363 100644
--- a/tests/MyGame/Example/Monster.py
+++ b/tests/MyGame/Example/Monster.py
@@ -15,7 +15,6 @@ class Monster(object):
         x.Init(buf, n + offset)
         return x
 
-
     # Monster
     def Init(self, buf, pos):
         self._tab = flatbuffers.table.Table(buf, pos)
diff --git a/tests/MyGame/Example/Stat.go b/tests/MyGame/Example/Stat.go
index 4ccca344..8c56bfb0 100644
--- a/tests/MyGame/Example/Stat.go
+++ b/tests/MyGame/Example/Stat.go
@@ -9,6 +9,13 @@ type Stat struct {
 	_tab flatbuffers.Table
 }
 
+func GetRootAsStat(buf []byte, offset flatbuffers.UOffsetT) *Stat {
+	n := flatbuffers.GetUOffsetT(buf[offset:])
+	x := &Stat{}
+	x.Init(buf, n + offset)
+	return x
+}
+
 func (rcv *Stat) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i
diff --git a/tests/MyGame/Example/Stat.py b/tests/MyGame/Example/Stat.py
index 71577e86..b0e251d4 100644
--- a/tests/MyGame/Example/Stat.py
+++ b/tests/MyGame/Example/Stat.py
@@ -7,6 +7,13 @@ import flatbuffers
 class Stat(object):
     __slots__ = ['_tab']
 
+    @classmethod
+    def GetRootAsStat(cls, buf, offset):
+        n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
+        x = Stat()
+        x.Init(buf, n + offset)
+        return x
+
     # Stat
     def Init(self, buf, pos):
         self._tab = flatbuffers.table.Table(buf, pos)
diff --git a/tests/MyGame/Example/TestSimpleTableWithEnum.go b/tests/MyGame/Example/TestSimpleTableWithEnum.go
index d59fa6b1..96218e58 100644
--- a/tests/MyGame/Example/TestSimpleTableWithEnum.go
+++ b/tests/MyGame/Example/TestSimpleTableWithEnum.go
@@ -9,6 +9,13 @@ type TestSimpleTableWithEnum struct {
 	_tab flatbuffers.Table
 }
 
+func GetRootAsTestSimpleTableWithEnum(buf []byte, offset flatbuffers.UOffsetT) *TestSimpleTableWithEnum {
+	n := flatbuffers.GetUOffsetT(buf[offset:])
+	x := &TestSimpleTableWithEnum{}
+	x.Init(buf, n + offset)
+	return x
+}
+
 func (rcv *TestSimpleTableWithEnum) Init(buf []byte, i flatbuffers.UOffsetT) {
 	rcv._tab.Bytes = buf
 	rcv._tab.Pos = i
diff --git a/tests/MyGame/Example/TestSimpleTableWithEnum.py b/tests/MyGame/Example/TestSimpleTableWithEnum.py
index 5b950979..8d64e971 100644
--- a/tests/MyGame/Example/TestSimpleTableWithEnum.py
+++ b/tests/MyGame/Example/TestSimpleTableWithEnum.py
@@ -7,6 +7,13 @@ import flatbuffers
 class TestSimpleTableWithEnum(object):
     __slots__ = ['_tab']
 
+    @classmethod
+    def GetRootAsTestSimpleTableWithEnum(cls, buf, offset):
+        n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
+        x = TestSimpleTableWithEnum()
+        x.Init(buf, n + offset)
+        return x
+
     # TestSimpleTableWithEnum
     def Init(self, buf, pos):
         self._tab = flatbuffers.table.Table(buf, pos)
diff --git a/tests/go_test.go b/tests/go_test.go
index b816eb29..e4e96de8 100644
--- a/tests/go_test.go
+++ b/tests/go_test.go
@@ -76,6 +76,9 @@ func TestAll(t *testing.T) {
 	CheckStructIsNotInlineError(t.Fatalf)
 	CheckFinishedBytesError(t.Fatalf)
 
+	// Verify that GetRootAs works for non-root tables
+	CheckGetRootAsForNonRootTable(t.Fatalf)
+
 	// Verify that using the generated Go code builds a buffer without
 	// returning errors:
 	generated, off := CheckGeneratedBuild(t.Fatalf)
@@ -1081,6 +1084,31 @@ func CheckManualBuild(fail func(string, ...interface{})) ([]byte, flatbuffers.UO
 	return b.Bytes, b.Head()
 }
 
+func CheckGetRootAsForNonRootTable(fail func(string, ...interface{})) {
+	b := flatbuffers.NewBuilder(0)
+	str := b.CreateString("MyStat")
+	example.StatStart(b)
+	example.StatAddId(b, str)
+	example.StatAddVal(b, 12345678)
+	example.StatAddCount(b, 12345)
+	stat_end := example.StatEnd(b)
+	b.Finish(stat_end)
+
+	stat := example.GetRootAsStat(b.Bytes, b.Head())
+
+	if got := stat.Id(); !bytes.Equal([]byte("MyStat"), got) {
+		fail(FailString("stat.Id()", "MyStat", got))
+	}
+
+	if got := stat.Val(); 12345678 != got {
+		fail(FailString("stat.Val()", 12345678, got))
+	}
+
+	if got := stat.Count(); 12345 != got {
+		fail(FailString("stat.Count()", 12345, got))
+	}
+}
+
 // CheckGeneratedBuild uses generated code to build the example Monster.
 func CheckGeneratedBuild(fail func(string, ...interface{})) ([]byte, flatbuffers.UOffsetT) {
 	b := flatbuffers.NewBuilder(0)
diff --git a/tests/py_test.py b/tests/py_test.py
index 82f83900..e6609449 100644
--- a/tests/py_test.py
+++ b/tests/py_test.py
@@ -1034,6 +1034,23 @@ class TestAllCodePathsOfExampleSchema(unittest.TestCase):
         self.assertEqual(7, mon2.Testhashs64Fnv1a())
         self.assertEqual(8, mon2.Testhashu64Fnv1a())
 
+    def test_getrootas_for_nonroot_table(self):
+        b = flatbuffers.Builder(0)
+        string = b.CreateString("MyStat")
+
+        MyGame.Example.Stat.StatStart(b)
+        MyGame.Example.Stat.StatAddId(b, string)
+        MyGame.Example.Stat.StatAddVal(b, 12345678)
+        MyGame.Example.Stat.StatAddCount(b, 12345)
+        stat = MyGame.Example.Stat.StatEnd(b)
+        b.Finish(stat)
+
+        stat2 = MyGame.Example.Stat.Stat.GetRootAsStat(b.Bytes, b.Head())
+
+        self.assertEqual(b"MyStat", stat2.Id())
+        self.assertEqual(12345678, stat2.Val())
+        self.assertEqual(12345, stat2.Count())
+
 
 class TestVtableDeduplication(unittest.TestCase):
     ''' TestVtableDeduplication verifies that vtables are deduplicated. '''