Resolve #4530: Added vector union support for Python

Author: surculus12

scripts/generate_code.py

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

src/idl_gen_python.cpp

@@ -1076,6 +1076,39 @@ 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";
+
+    // TODO(rw): this works and is not the good way to it:
+    bool is_native_table = TypeName(field) == "*flatbuffers.Table";
+    if (is_native_table) {
+    code +=
+    Indent + Indent + Indent + "from flatbuffers.table import Table\n";
+    } else if (parser_.opts.include_dependence_headers) {
+    code += Indent + Indent + Indent;
+    code += "from " + GenPackageReference(field.value.type) + " import " +
+    TypeName(field) + "\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,
@@ -1521,6 +1554,8 @@ class PythonGenerator : public BaseGenerator {
           auto vectortype = field.value.type.VectorType();
           if (vectortype.base_type == BASE_TYPE_STRUCT) {
             GetMemberOfVectorOfStruct(struct_def, field, code_ptr, imports);
+          } else if (vectortype.base_type == BASE_TYPE_UNION) {
+            GetMemberOfVectorOfUnion(struct_def, field, code_ptr);
           } else {
             GetMemberOfVectorOfNonStruct(struct_def, field, code_ptr);
             if (parser_.opts.python_gen_numpy) {
@@ -1780,6 +1815,9 @@ class PythonGenerator : public BaseGenerator {
         import_list->insert("import " + package_reference);
       }
       field_type = "Optional[List[" + field_type + "]";
+    } else if (base_type == BASE_TYPE_UNION) {
+      GenUnionInit(field, field_type_ptr, import_list, import_typing_list);
+      field_type = "Optional[List[" + field_type + "]]";
     } else {
       field_type = "Optional[List[" +
                    GetBasePythonTypeForScalarAndString(base_type) + "]]";

src/idl_parser.cpp

@@ -2702,7 +2702,7 @@ bool Parser::SupportsAdvancedUnionFeatures() const {
   return (opts.lang_to_generate &
           ~(IDLOptions::kCpp | IDLOptions::kTs | IDLOptions::kPhp |
             IDLOptions::kJava | IDLOptions::kCSharp | IDLOptions::kKotlin |
-            IDLOptions::kBinary | IDLOptions::kSwift | IDLOptions::kNim |
+            IDLOptions::kBinary | IDLOptions::kSwift | IDLOptions::kPython | IDLOptions::kNim |
             IDLOptions::kJson | IDLOptions::kKotlinKmp)) == 0;
 }
 

tests/py_test.py

@@ -54,6 +54,10 @@
 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
 
 
 def create_namespace_shortcut(is_onefile):
@@ -301,6 +305,51 @@ 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)
+
+    union_vector.Attacker.Start(b)
+    union_vector.Attacker.AddSwordAttackDamage(b, 1)
+    attacker_offset = union_vector.Attacker.End(b)
+
+    union_vector.Attacker.Start(b)
+    union_vector.Attacker.AddSwordAttackDamage(b, 2)
+    attacker_offset2 = union_vector.Attacker.End(b)
+
+    characters = [attacker_offset, attacker_offset2]
+    character_types = [union_vector.Character.Character.MuLan, union_vector.Character.Character.MuLan]
+
+    union_vector.Movie.StartCharactersTypeVector(b, len(character_types))
+    for character_type in reversed(character_types):
+      b.PrependByte(character_type)
+    character_types_offset = b.EndVector()
+
+    union_vector.Movie.StartCharactersVector(b, len(characters))
+    for character in reversed(characters):
+      b.PrependUOffsetTRelative(character)
+    characters_offset = b.EndVector()
+
+    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)
+
+    buf = b.Output()
+    movie = union_vector.Movie.Movie.GetRootAsMovie(buf, 0)
+
+    self.assertEqual(movie.CharactersTypeLength(), len(character_types))
+    self.assertEqual(movie.CharactersLength(), len(characters))
+    self.assertEqual(movie.CharactersType(0), character_types[0])
+
+    character = union_vector.Attacker.Attacker()
+    character.Init(movie.Characters(0).Bytes, movie.Characters(0).Pos)
+    self.assertEqual(character.SwordAttackDamage(), 1)
+    character.Init(movie.Characters(1).Bytes, movie.Characters(1).Pos)
+    self.assertEqual(character.SwordAttackDamage(), 2)
+
   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.String(table.Pos)
+        return union
+    if unionType == Character().Unused:
+        tab = Table(table.Bytes, table.Pos)
+        union = tab.String(table.Pos)
+        return union
+    return None

tests/union_vector/FallingTub.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 FallingTub(object):
+    __slots__ = ['_tab']
+
+    @classmethod
+    def SizeOf(cls):
+        return 4
+
+    # FallingTub
+    def Init(self, buf, pos):
+        self._tab = flatbuffers.table.Table(buf, pos)
+
+    # FallingTub
+    def Weight(self): return self._tab.Get(flatbuffers.number_types.Int32Flags, self._tab.Pos + flatbuffers.number_types.UOffsetTFlags.py_type(0))
+
+def CreateFallingTub(builder, weight):
+    builder.Prep(4, 4)
+    builder.PrependInt32(weight)
+    return builder.Offset()
+
+
+class FallingTubT(object):
+
+    # FallingTubT
+    def __init__(self):
+        self.weight = 0  # type: int
+
+    @classmethod
+    def InitFromBuf(cls, buf, pos):
+        fallingTub = FallingTub()
+        fallingTub.Init(buf, pos)
+        return cls.InitFromObj(fallingTub)
+
+    @classmethod
+    def InitFromObj(cls, fallingTub):
+        x = FallingTubT()
+        x._UnPack(fallingTub)
+        return x
+
+    # FallingTubT
+    def _UnPack(self, fallingTub):
+        if fallingTub is None:
+            return
+        self.weight = fallingTub.Weight()
+
+    # FallingTubT
+    def Pack(self, builder):
+        return CreateFallingTub(builder, self.weight)

tests/union_vector/Gadget.py

@@ -0,0 +1,20 @@
+# automatically generated by the FlatBuffers compiler, do not modify
+
+# namespace: 
+
+class Gadget(object):
+    NONE = 0
+    FallingTub = 1
+    HandFan = 2
+
+def GadgetCreator(unionType, table):
+    from flatbuffers.table import Table
+    if not isinstance(table, Table):
+        return None
+    if unionType == Gadget().FallingTub:
+        import FallingTub
+        return FallingTub.FallingTubT.InitFromBuf(table.Bytes, table.Pos)
+    if unionType == Gadget().HandFan:
+        import HandFan
+        return HandFan.HandFanT.InitFromBuf(table.Bytes, table.Pos)
+    return None