Merge pull request #344 from rw/py-builder-ergonomics

Python: Improve Builder user interface.

python/flatbuffers/builder.py

@@ -32,7 +32,7 @@ class OffsetArithmeticError(RuntimeError):
     pass
 
 
-class NotInObjectError(RuntimeError):
+class IsNotNestedError(RuntimeError):
     """
     Error caused by using a Builder to write Object data when not inside
     an Object.
@@ -40,7 +40,7 @@ class NotInObjectError(RuntimeError):
     pass
 
 
-class ObjectIsNestedError(RuntimeError):
+class IsNestedError(RuntimeError):
     """
     Error caused by using a Builder to begin an Object when an Object is
     already being built.
@@ -63,6 +63,12 @@ class BuilderSizeError(RuntimeError):
     """
     pass
 
+class BuilderNotFinishedError(RuntimeError):
+    """
+    Error caused by not calling `Finish` before calling `Output`.
+    """
+    pass
+
 
 # VtableMetadataFields is the count of metadata fields in each vtable.
 VtableMetadataFields = 2
@@ -85,7 +91,7 @@ class Builder(object):
     """
 
     __slots__ = ("Bytes", "current_vtable", "head", "minalign", "objectEnd",
-                 "vtables")
+                 "vtables", "nested", "finished")
 
     """
     Maximum buffer size constant, in bytes.
@@ -110,13 +116,19 @@ class Builder(object):
         self.minalign = 1
         self.objectEnd = None
         self.vtables = []
+        self.nested = False
+        self.finished = False
 
     def Output(self):
         """
         Output returns the portion of the buffer that has been used for
-        writing data.
+        writing data. It raises BuilderNotFinishedError if the buffer has not
+        been finished with `Finish`.
         """
 
+        if not self.finished:
+            raise BuilderNotFinishedError()
+
         return self.Bytes[self.Head():]
 
     def StartObject(self, numfields):
@@ -128,6 +140,7 @@ class Builder(object):
         self.current_vtable = [0 for _ in range_func(numfields)]
         self.objectEnd = self.Offset()
         self.minalign = 1
+        self.nested = True
 
     def WriteVtable(self):
         """
@@ -236,10 +249,8 @@ class Builder(object):
 
     def EndObject(self):
         """EndObject writes data necessary to finish object construction."""
-        if self.current_vtable is None:
-            msg = ("flatbuffers: Tried to write the end of an Object when "
-                   "the Builder was not currently writing an Object.")
-            raise NotInObjectError(msg)
+        self.assertNested()
+        self.nested = False
         return self.WriteVtable()
 
     def growByteBuffer(self):
@@ -336,6 +347,7 @@ class Builder(object):
         """
 
         self.assertNotNested()
+        self.nested = True
         self.Prep(N.Uint32Flags.bytewidth, elemSize*numElems)
         self.Prep(alignment, elemSize*numElems)  # In case alignment > int.
         return self.Offset()
@@ -343,6 +355,8 @@ class Builder(object):
     def EndVector(self, vectorNumElems):
         """EndVector writes data necessary to finish vector construction."""
 
+        self.assertNested()
+        self.nested = False
         # we already made space for this, so write without PrependUint32
         self.PlaceUOffsetT(vectorNumElems)
         return self.Offset()
@@ -351,6 +365,7 @@ class Builder(object):
         """CreateString writes a null-terminated byte string as a vector."""
 
         self.assertNotNested()
+        self.nested = True
 
         if isinstance(s, compat.string_types):
             x = s.encode()
@@ -369,18 +384,24 @@ class Builder(object):
 
         return self.EndVector(len(x))
 
+    def assertNested(self):
+        """
+        Check that we are in the process of building an object.
+        """
+
+        if not self.nested:
+            raise IsNotNestedError()
+
     def assertNotNested(self):
         """
         Check that no other objects are being built while making this
         object. If not, raise an exception.
         """
 
-        if self.current_vtable is not None:
-            msg = ("flatbuffers: Tried to write a new Object when the "
-                   "Builder was already writing an Object.")
-            raise ObjectIsNestedError(msg)
+        if self.nested:
+            raise IsNestedError()
 
-    def assertNested(self, obj):
+    def assertStructIsInline(self, obj):
         """
         Structs are always stored inline, so need to be created right
         where they are used. You'll get this error if you created it
@@ -399,11 +420,7 @@ class Builder(object):
         buffer.
 
         """
-        if self.current_vtable is None:
-            msg = ("flatbuffers: Tried to write an Object field when "
-                   "the Builder was not currently writing an Object.")
-            raise NotInObjectError(msg)
-
+        self.assertNested()
         self.current_vtable[slotnum] = self.Offset()
 
     def Finish(self, rootTable):
@@ -411,6 +428,7 @@ class Builder(object):
         N.enforce_number(rootTable, N.UOffsetTFlags)
         self.Prep(self.minalign, N.UOffsetTFlags.bytewidth)
         self.PrependUOffsetTRelative(rootTable)
+        self.finished = True
         return self.Head()
 
     def Prepend(self, flags, off):
@@ -470,7 +488,7 @@ class Builder(object):
 
         N.enforce_number(d, N.UOffsetTFlags)
         if x != d:
-            self.assertNested(x)
+            self.assertStructIsInline(x)
             self.Slot(v)
 
     def PrependBool(self, x): self.Prepend(N.BoolFlags, x)

tests/py_test.py

@@ -309,7 +309,7 @@ class TestByteLayout(unittest.TestCase):
 
         want_ints = list(map(integerize, want_chars_or_ints))
         want = bytearray(want_ints)
-        got = builder.Output()
+        got = builder.Bytes[builder.Head():] # use the buffer directly
         self.assertEqual(want, got)
 
     def test_numbers(self):
@@ -878,24 +878,23 @@ class TestAllCodePathsOfExampleSchema(unittest.TestCase):
         b = flatbuffers.Builder(0)
 
         # make a child Monster within a vector of Monsters:
-        MyGame.Example.Monster.MonsterStartTestarrayoftablesVector(b, 1)
-
         MyGame.Example.Monster.MonsterStart(b)
         MyGame.Example.Monster.MonsterAddHp(b, 99)
         sub_monster = MyGame.Example.Monster.MonsterEnd(b)
-        b.Finish(sub_monster)
 
-        tables = b.EndVector(1)
+        # build the vector:
+        MyGame.Example.Monster.MonsterStartTestarrayoftablesVector(b, 1)
+        b.PrependUOffsetTRelative(sub_monster)
+        vec = b.EndVector(1)
 
         # make the parent monster and include the vector of Monster:
         MyGame.Example.Monster.MonsterStart(b)
-        MyGame.Example.Monster.MonsterAddTestarrayoftables(b, tables)
+        MyGame.Example.Monster.MonsterAddTestarrayoftables(b, vec)
         mon = MyGame.Example.Monster.MonsterEnd(b)
         b.Finish(mon)
 
         # inspect the resulting data:
-        mon2 = MyGame.Example.Monster.Monster.GetRootAsMonster(b.Bytes,
-                                                               b.Head())
+        mon2 = MyGame.Example.Monster.Monster.GetRootAsMonster(b.Output(), 0)
         self.assertEqual(99, mon2.Testarrayoftables(0).Hp())
         self.assertEqual(1, mon2.TestarrayoftablesLength())
 
@@ -1050,7 +1049,7 @@ class TestVtableDeduplication(unittest.TestCase):
         b.PrependInt16Slot(3, 99, 0)
         obj2 = b.EndObject()
 
-        got = b.Output()
+        got = b.Bytes[b.Head():]
 
         want = bytearray([
             240, 255, 255, 255,  # == -12. offset to dedupped vtable.
@@ -1107,17 +1106,16 @@ class TestVtableDeduplication(unittest.TestCase):
 
 
 class TestExceptions(unittest.TestCase):
-    def test_not_in_object_error(self):
-        b = flatbuffers.Builder(0)
-        exc = None
-        assertRaises(self, lambda: b.EndObject(),
-                     flatbuffers.builder.NotInObjectError)
-
     def test_object_is_nested_error(self):
         b = flatbuffers.Builder(0)
         b.StartObject(0)
         assertRaises(self, lambda: b.StartObject(0),
-                     flatbuffers.builder.ObjectIsNestedError)
+                     flatbuffers.builder.IsNestedError)
+
+    def test_object_is_not_nested_error(self):
+        b = flatbuffers.Builder(0)
+        assertRaises(self, lambda: b.EndObject(),
+                     flatbuffers.builder.IsNotNestedError)
 
     def test_struct_is_not_inline_error(self):
         b = flatbuffers.Builder(0)
@@ -1135,7 +1133,12 @@ class TestExceptions(unittest.TestCase):
         b.StartObject(0)
         s = 'test1'
         assertRaises(self, lambda: b.CreateString(s),
-                     flatbuffers.builder.ObjectIsNestedError)
+                     flatbuffers.builder.IsNestedError)
+
+    def test_finished_bytes_error(self):
+        b = flatbuffers.Builder(0)
+        assertRaises(self, lambda: b.Output(),
+                     flatbuffers.builder.BuilderNotFinishedError)
 
 
 def CheckAgainstGoldDataGo():