From e9d29c21a70c91d6e33c7573fd6beead01bdbe7e Mon Sep 17 00:00:00 2001
From: Bharat Tak <bharat.tak@wingtra.com>
Date: Fri, 18 Oct 2019 00:25:05 +0200
Subject: [PATCH] Python: Add forceDefaults opt to python Builder (#5564)

* Add forceDefaults opt to python Builder

* Add test functions for force_default option for python builder

* Simplify

* Add force default test for UOffsetTFlags
---
 python/flatbuffers/builder.py | 17 +++++++++++++----
 tests/py_test.py              | 35 +++++++++++++++++++++++++++++++++++
 2 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/python/flatbuffers/builder.py b/python/flatbuffers/builder.py
index d04ee854d..478cdcd5e 100644
--- a/python/flatbuffers/builder.py
+++ b/python/flatbuffers/builder.py
@@ -103,7 +103,7 @@ class Builder(object):
 
     ## @cond FLATBUFFERS_INTENRAL
     __slots__ = ("Bytes", "current_vtable", "head", "minalign", "objectEnd",
-                 "vtables", "nested", "finished")
+                 "vtables", "nested", "forceDefaults", "finished")
 
     """Maximum buffer size constant, in bytes.
 
@@ -131,10 +131,10 @@ class Builder(object):
         self.objectEnd = None
         self.vtables = {}
         self.nested = False
+        self.forceDefaults = False
         ## @endcond
         self.finished = False
 
-
     def Output(self):
         """Return the portion of the buffer that has been used for writing data.
 
@@ -552,7 +552,7 @@ class Builder(object):
     def PrependSlot(self, flags, o, x, d):
         N.enforce_number(x, flags)
         N.enforce_number(d, flags)
-        if x != d:
+        if x != d or self.forceDefaults:
             self.Prepend(flags, x)
             self.Slot(o)
 
@@ -589,7 +589,7 @@ class Builder(object):
         be set to zero and no other data will be written.
         """
 
-        if x != d:
+        if x != d or self.forceDefaults:
             self.PrependUOffsetTRelative(x)
             self.Slot(o)
 
@@ -691,6 +691,15 @@ class Builder(object):
         """
         self.Prepend(N.Float64Flags, x)
 
+    def ForceDefaults(self, forceDefaults):
+        """
+        In order to save space, fields that are set to their default value
+        don't get serialized into the buffer. Forcing defaults provides a
+        way to manually disable this optimization. When set to `True`, will
+        always serialize default values.
+        """
+        self.forceDefaults = forceDefaults
+
 ##############################################################
 
     ## @cond FLATBUFFERS_INTERNAL
diff --git a/tests/py_test.py b/tests/py_test.py
index 0a661f3d9..4550ac7da 100644
--- a/tests/py_test.py
+++ b/tests/py_test.py
@@ -1159,6 +1159,41 @@ def make_monster_from_generated_code(sizePrefix = False, file_identifier=None):
     return b.Bytes, b.Head()
 
 
+class TestBuilderForceDefaults(unittest.TestCase):
+    """Verify that the builder adds default values when forced."""
+
+    test_flags = [N.BoolFlags(), N.Uint8Flags(), N.Uint16Flags(), \
+                  N.Uint32Flags(), N.Uint64Flags(), N.Int8Flags(), \
+                  N.Int16Flags(), N.Int32Flags(), N.Int64Flags(), \
+                  N.Float32Flags(), N.Float64Flags(), N.UOffsetTFlags()]
+    def test_default_force_defaults(self):
+        for flag in self.test_flags:
+            b = flatbuffers.Builder(0)
+            b.StartObject(1)
+            stored_offset = b.Offset()
+            if flag != N.UOffsetTFlags():
+                b.PrependSlot(flag, 0, 0, 0)
+            else:
+                b.PrependUOffsetTRelativeSlot(0, 0, 0)
+            end_offset = b.Offset()
+            b.EndObject()
+            self.assertEqual(0, end_offset - stored_offset)
+
+    def test_force_defaults_true(self):
+        for flag in self.test_flags:
+            b = flatbuffers.Builder(0)
+            b.ForceDefaults(True)
+            b.StartObject(1)
+            stored_offset = b.Offset()
+            if flag != N.UOffsetTFlags():
+                b.PrependSlot(flag, 0, 0, 0)
+            else:
+                b.PrependUOffsetTRelativeSlot(0, 0, 0)
+            end_offset = b.Offset()
+            b.EndObject()
+            self.assertEqual(flag.bytewidth, end_offset - stored_offset)
+
+
 class TestAllCodePathsOfExampleSchema(unittest.TestCase):
     def setUp(self, *args, **kwargs):
         super(TestAllCodePathsOfExampleSchema, self).setUp(*args, **kwargs)