Added vector union support for Python

python/flatbuffers/table.py

@@ -53,6 +53,13 @@ def String(self, off):
         length = encode.Get(N.UOffsetTFlags.packer_type, self.Bytes, off)
         return bytes(self.Bytes[start:start+length])
 
+    def UnionString(self, off):
+        """String gets a string from data stored inside the flatbuffer."""
+        N.enforce_number(off, N.UOffsetTFlags)
+        start = off + N.UOffsetTFlags.bytewidth
+        length = encode.Get(N.UOffsetTFlags.packer_type, self.Bytes, off)
+        return bytes(self.Bytes[start:start+length])
+
     def VectorLen(self, off):
         """VectorLen retrieves the length of the vector whose offset is stored
            at "off" in this object."""

scripts/generate_code.py

@@ -234,7 +234,7 @@ def glob(path, pattern):
 )
 
 flatc(
-    BASE_OPTS + CPP_OPTS + CS_OPTS + TS_OPTS + JAVA_OPTS + KOTLIN_OPTS + PHP_OPTS,
+    BASE_OPTS + CPP_OPTS + CS_OPTS + TS_OPTS + JAVA_OPTS + KOTLIN_OPTS + PHP_OPTS + PYTHON_OPTS,
     prefix="union_vector",
     schema="union_vector/union_vector.fbs",
 )

src/idl_gen_python.cpp

@@ -360,6 +360,30 @@ class PythonGenerator : public BaseGenerator {
     code += Indent + Indent + "return None\n\n";
   }
 
+    // Get the value of a vector's union member.
+    void GetMemberOfVectorOfUnion(const StructDef &struct_def,
+                                 const FieldDef &field,
+                                 std::string *code_ptr) const {
+      auto &code = *code_ptr;
+      auto vectortype = field.value.type.VectorType();
+
+      GenReceiver(struct_def, code_ptr);
+      code += namer_.Method(field);
+      code += "(self, j):" + OffsetPrefix(field);
+      code += Indent + Indent + Indent + "x = self._tab.Vector(o)\n";
+      code += Indent + Indent + Indent;
+      code += "x += flatbuffers.number_types.UOffsetTFlags.py_type(j) * ";
+      code += NumToString(InlineSize(vectortype)) + "\n";
+      code += Indent + Indent + Indent;
+      code += "x -= self._tab.Pos\n";
+      code +=
+          Indent + Indent + Indent + "from flatbuffers.table import Table\n";
+      code += Indent + Indent + Indent + "obj = Table(bytearray(), 0)\n";
+      code += Indent + Indent + Indent + GenGetter(field.value.type);
+      code += "obj, x)\n" + Indent + Indent + Indent + "return obj\n";
+      code += Indent + Indent + "return None\n\n";
+    }
+
   // Get the value of a vector's non-struct member. Uses a named return
   // argument to conveniently set the zero value for the result.
   void GetMemberOfVectorOfNonStruct(const StructDef &struct_def,
@@ -728,6 +752,8 @@ class PythonGenerator : public BaseGenerator {
           auto vectortype = field.value.type.VectorType();
           if (vectortype.base_type == BASE_TYPE_STRUCT) {
             GetMemberOfVectorOfStruct(struct_def, field, code_ptr);
+          } else if (vectortype.base_type == BASE_TYPE_UNION) {
+            GetMemberOfVectorOfUnion(struct_def, field, code_ptr);
           } else {
             GetMemberOfVectorOfNonStruct(struct_def, field, code_ptr);
             GetVectorOfNonStructAsNumpy(struct_def, field, code_ptr);
@@ -964,6 +990,9 @@ class PythonGenerator : public BaseGenerator {
         import_list->insert("import " + package_reference);
       }
       field_type = "List[" + field_type;
+    } else if (base_type == BASE_TYPE_UNION) {
+      GenUnionInit(field, field_type_ptr, import_list, import_typing_list);
+      field_type = "List[" + field_type + "]";
     } else {
       field_type =
           "List[" + GetBasePythonTypeForScalarAndString(base_type) + "]";
@@ -1579,7 +1608,7 @@ class PythonGenerator : public BaseGenerator {
     code +=
         GenIndents(1) + "if unionType == " + union_type + "()." + variant + ":";
     code += GenIndents(2) + "tab = Table(table.Bytes, table.Pos)";
-    code += GenIndents(2) + "union = tab.String(table.Pos)";
+    code += GenIndents(2) + "union = tab.UnionString(table.Pos)";
     code += GenIndents(2) + "return union";
   }
 

src/idl_parser.cpp

@@ -2496,10 +2496,14 @@ bool Parser::SupportsDefaultVectorsAndStrings() const {
 }
 
 bool Parser::SupportsAdvancedUnionFeatures() const {
+  // Advanced Union Features refers to the following features:
+  //    - Vectors of unions
+  //    - Strings in unions
+  //    - structs in unions
   return (opts.lang_to_generate &
           ~(IDLOptions::kCpp | IDLOptions::kTs | IDLOptions::kPhp |
             IDLOptions::kJava | IDLOptions::kCSharp | IDLOptions::kKotlin |
-            IDLOptions::kBinary | IDLOptions::kSwift)) == 0;
+            IDLOptions::kBinary | IDLOptions::kSwift | IDLOptions::kPython)) == 0;
 }
 
 bool Parser::SupportsAdvancedArrayFeatures() const {

tests/py_test.py

@@ -50,6 +50,12 @@
 import monster_test_generated  # the one-file version
 import optional_scalars
 import optional_scalars.ScalarStuff
+import union_vector
+import union_vector.Attacker
+import union_vector.Character
+import union_vector.Movie
+import union_vector.Rapunzel
+import union_vector.BookReader
 
 
 def create_namespace_shortcut(is_onefile):
@@ -276,6 +282,86 @@ def test_default_values_with_pack_and_unpack(self):
     self.assertEqual(monster2.VectorOfEnumsLength(), 0)
     self.assertTrue(monster2.VectorOfEnumsIsNone())
 
+  def test_union_vectors_with_pack_and_unpack(self):
+    b = flatbuffers.Builder(0)
+
+    # Types of the characters
+    types = [
+      union_vector.Character.Character.MuLan,
+      union_vector.Character.Character.Rapunzel,
+      union_vector.Character.Character.Belle,
+      union_vector.Character.Character.BookFan,
+      union_vector.Character.Character.Other,
+      union_vector.Character.Character.Unused,
+    ]
+    # Pack the attacker manually
+    union_vector.Attacker.Start(b)
+    union_vector.Attacker.AddSwordAttackDamage(b, 1)
+    attacker_offset = union_vector.Attacker.End(b)
+    # Prepare the rest of the characters
+    characters = [
+      attacker_offset,
+      union_vector.Rapunzel.CreateRapunzel(b, 2),
+      union_vector.BookReader.CreateBookReader(b, 3),
+      union_vector.BookReader.CreateBookReader(b, 4),
+      b.CreateString("Other", "utf-8"),
+      b.CreateString("Unused", "utf-8")
+    ]
+
+    # Pack the types
+    union_vector.Movie.StartCharactersTypeVector(b, len(types))
+    for character_type in reversed(types):
+      b.PrependByte(character_type)
+    character_types_offset = b.EndVector()
+    # Pack the characters
+    union_vector.Movie.StartCharactersVector(b, len(characters))
+    for character in reversed(characters):
+      b.PrependUOffsetTRelative(character)
+    characters_offset = b.EndVector()
+
+    # Pack the movie object
+    union_vector.Movie.Start(b)
+    union_vector.Movie.AddMainCharacterType(b, 0)
+    union_vector.Movie.AddMainCharacter(b, 0)
+    union_vector.Movie.AddCharactersType(b, character_types_offset)
+    union_vector.Movie.AddCharacters(b, characters_offset)
+    movie_offset = union_vector.Movie.End(b)
+    b.Finish(movie_offset)
+
+    # Unpack the movie object
+    buf = b.Output()
+    movie = union_vector.Movie.Movie.GetRootAsMovie(buf, 0)
+
+    self.assertEqual(movie.CharactersType(0), union_vector.Character.Character.MuLan)
+    self.assertEqual(movie.CharactersType(1), union_vector.Character.Character.Rapunzel)
+    self.assertEqual(movie.CharactersType(2), union_vector.Character.Character.Belle)
+    self.assertEqual(movie.CharactersType(3), union_vector.Character.Character.BookFan)
+    self.assertEqual(movie.CharactersType(4), union_vector.Character.Character.Other)
+    self.assertEqual(movie.CharactersType(5), union_vector.Character.Character.Unused)
+
+    attacker = union_vector.Attacker.Attacker()
+    attacker.Init(movie.Characters(0).Bytes, movie.Characters(0).Pos)
+    self.assertEqual(attacker.SwordAttackDamage(), 1)
+    rapunzel = union_vector.Rapunzel.Rapunzel()
+    rapunzel.Init(movie.Characters(1).Bytes, movie.Characters(1).Pos)
+    self.assertEqual(rapunzel.HairLength(), 2)
+    book_reader = union_vector.BookReader.BookReader()
+    book_reader.Init(movie.Characters(2).Bytes, movie.Characters(2).Pos)
+    self.assertEqual(book_reader.BooksRead(), 3)
+    book_reader.Init(movie.Characters(3).Bytes, movie.Characters(3).Pos)
+    self.assertEqual(book_reader.BooksRead(), 4)
+
+    other = union_vector.Character.CharacterCreator(
+      union_vector.Character.Character.Other,
+      movie.Characters(4)
+    )
+    self.assertEqual(other.decode("utf-8"), "Other")
+    unused = union_vector.Character.CharacterCreator(
+      union_vector.Character.Character.Unused,
+      movie.Characters(5)
+    )
+    self.assertEqual(unused.decode("utf-8"), "Unused")
+
   def test_optional_scalars_with_pack_and_unpack(self):
     """ Serializes and deserializes between a buffer with optional values (no
         specific values are filled when the buffer is created) and its python

tests/union_vector/Attacker.py

@@ -0,0 +1,77 @@
+# automatically generated by the FlatBuffers compiler, do not modify
+
+# namespace: 
+
+import flatbuffers
+from flatbuffers.compat import import_numpy
+np = import_numpy()
+
+class Attacker(object):
+    __slots__ = ['_tab']
+
+    @classmethod
+    def GetRootAs(cls, buf, offset=0):
+        n = flatbuffers.encode.Get(flatbuffers.packer.uoffset, buf, offset)
+        x = Attacker()
+        x.Init(buf, n + offset)
+        return x
+
+    @classmethod
+    def GetRootAsAttacker(cls, buf, offset=0):
+        """This method is deprecated. Please switch to GetRootAs."""
+        return cls.GetRootAs(buf, offset)
+    @classmethod
+    def AttackerBufferHasIdentifier(cls, buf, offset, size_prefixed=False):
+        return flatbuffers.util.BufferHasIdentifier(buf, offset, b"\x4D\x4F\x56\x49", size_prefixed=size_prefixed)
+
+    # Attacker
+    def Init(self, buf, pos):
+        self._tab = flatbuffers.table.Table(buf, pos)
+
+    # Attacker
+    def SwordAttackDamage(self):
+        o = flatbuffers.number_types.UOffsetTFlags.py_type(self._tab.Offset(4))
+        if o != 0:
+            return self._tab.Get(flatbuffers.number_types.Int32Flags, o + self._tab.Pos)
+        return 0
+
+def AttackerStart(builder): builder.StartObject(1)
+def Start(builder):
+    return AttackerStart(builder)
+def AttackerAddSwordAttackDamage(builder, swordAttackDamage): builder.PrependInt32Slot(0, swordAttackDamage, 0)
+def AddSwordAttackDamage(builder, swordAttackDamage):
+    return AttackerAddSwordAttackDamage(builder, swordAttackDamage)
+def AttackerEnd(builder): return builder.EndObject()
+def End(builder):
+    return AttackerEnd(builder)
+
+class AttackerT(object):
+
+    # AttackerT
+    def __init__(self):
+        self.swordAttackDamage = 0  # type: int
+
+    @classmethod
+    def InitFromBuf(cls, buf, pos):
+        attacker = Attacker()
+        attacker.Init(buf, pos)
+        return cls.InitFromObj(attacker)
+
+    @classmethod
+    def InitFromObj(cls, attacker):
+        x = AttackerT()
+        x._UnPack(attacker)
+        return x
+
+    # AttackerT
+    def _UnPack(self, attacker):
+        if attacker is None:
+            return
+        self.swordAttackDamage = attacker.SwordAttackDamage()
+
+    # AttackerT
+    def Pack(self, builder):
+        AttackerStart(builder)
+        AttackerAddSwordAttackDamage(builder, self.swordAttackDamage)
+        attacker = AttackerEnd(builder)
+        return attacker

tests/union_vector/BookReader.py

@@ -0,0 +1,55 @@
+# automatically generated by the FlatBuffers compiler, do not modify
+
+# namespace: 
+
+import flatbuffers
+from flatbuffers.compat import import_numpy
+np = import_numpy()
+
+class BookReader(object):
+    __slots__ = ['_tab']
+
+    @classmethod
+    def SizeOf(cls):
+        return 4
+
+    # BookReader
+    def Init(self, buf, pos):
+        self._tab = flatbuffers.table.Table(buf, pos)
+
+    # BookReader
+    def BooksRead(self): return self._tab.Get(flatbuffers.number_types.Int32Flags, self._tab.Pos + flatbuffers.number_types.UOffsetTFlags.py_type(0))
+
+def CreateBookReader(builder, booksRead):
+    builder.Prep(4, 4)
+    builder.PrependInt32(booksRead)
+    return builder.Offset()
+
+
+class BookReaderT(object):
+
+    # BookReaderT
+    def __init__(self):
+        self.booksRead = 0  # type: int
+
+    @classmethod
+    def InitFromBuf(cls, buf, pos):
+        bookReader = BookReader()
+        bookReader.Init(buf, pos)
+        return cls.InitFromObj(bookReader)
+
+    @classmethod
+    def InitFromObj(cls, bookReader):
+        x = BookReaderT()
+        x._UnPack(bookReader)
+        return x
+
+    # BookReaderT
+    def _UnPack(self, bookReader):
+        if bookReader is None:
+            return
+        self.booksRead = bookReader.BooksRead()
+
+    # BookReaderT
+    def Pack(self, builder):
+        return CreateBookReader(builder, self.booksRead)

tests/union_vector/Character.py

@@ -0,0 +1,38 @@
+# automatically generated by the FlatBuffers compiler, do not modify
+
+# namespace: 
+
+class Character(object):
+    NONE = 0
+    MuLan = 1
+    Rapunzel = 2
+    Belle = 3
+    BookFan = 4
+    Other = 5
+    Unused = 6
+
+def CharacterCreator(unionType, table):
+    from flatbuffers.table import Table
+    if not isinstance(table, Table):
+        return None
+    if unionType == Character().MuLan:
+        import Attacker
+        return Attacker.AttackerT.InitFromBuf(table.Bytes, table.Pos)
+    if unionType == Character().Rapunzel:
+        import Rapunzel
+        return Rapunzel.RapunzelT.InitFromBuf(table.Bytes, table.Pos)
+    if unionType == Character().Belle:
+        import BookReader
+        return BookReader.BookReaderT.InitFromBuf(table.Bytes, table.Pos)
+    if unionType == Character().BookFan:
+        import BookReader
+        return BookReader.BookReaderT.InitFromBuf(table.Bytes, table.Pos)
+    if unionType == Character().Other:
+        tab = Table(table.Bytes, table.Pos)
+        union = tab.UnionString(table.Pos)
+        return union
+    if unionType == Character().Unused:
+        tab = Table(table.Bytes, table.Pos)
+        union = tab.UnionString(table.Pos)
+        return union
+    return None