Add class ByteCountingOutputStream

ByteCountingOutputStream is an output string stream implementation
that only counts the bytes.  This is useful for sizing an encoding.
----Release Notes----
[]
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=119995430

Author: gildea

google/cloud/dataflow/coders/slow_stream.py

@@ -59,6 +59,33 @@ def get(self):
     return ''.join(self.data)
 
 
+class ByteCountingOutputStream(OutputStream):
+  """A pure Python implementation of stream.ByteCountingOutputStream."""
+
+  def __init__(self):
+    # Note that we don't actually use any of the data initialized by our super.
+    super(ByteCountingOutputStream, self).__init__()
+    self.count = 0
+
+  def write(self, byte_array, nested=False):
+    blen = len(byte_array)
+    if nested:
+      self.write_var_int64(blen)
+    self.count += blen
+
+  def write_byte(self, _):
+    self.count += 1
+
+  def get_count(self):
+    return self.count
+
+  def get(self):
+    raise NotImplementedError
+
+  def __str__(self):
+    return '<%s %s>' % (self.__class__.__name__, self.count)
+
+
 class InputStream(object):
   """A pure Python implementation of stream.InputStream."""
 

google/cloud/dataflow/coders/stream.pxd

@@ -32,6 +32,17 @@ cdef class OutputStream(object):
   cdef extend(self, size_t missing)
 
 
+cdef class ByteCountingOutputStream(OutputStream):
+  cdef size_t count
+
+  cpdef write(self, bytes b, bint nested=*)
+  cpdef write_byte(self, unsigned char val)
+  cpdef write_bigendian_int64(self, libc.stdint.int64_t val)
+  cpdef write_bigendian_int32(self, libc.stdint.int32_t val)
+  cpdef size_t get_count(self)
+  cpdef bytes get(self)
+
+
 cdef class InputStream(object):
   cdef size_t pos
   cdef bytes all

google/cloud/dataflow/coders/stream.pyx

@@ -94,6 +94,44 @@ cdef class OutputStream(object):
     self.data = <char*>libc.stdlib.realloc(self.data, self.size)
 
 
+cdef class ByteCountingOutputStream(OutputStream):
+  """An output string stream implementation that only counts the bytes.
+
+  This implementation counts the number of bytes it "writes" but
+  doesn't actually write them anyway.  Thus it has write() but not
+  get().  get_count() returns how many bytes were written.
+
+  This is useful for sizing an encoding.
+  """
+
+  def __cinit__(self):
+    self.count = 0
+
+  cpdef write(self, bytes b, bint nested=False):
+    cdef size_t blen = len(b)
+    if nested:
+      self.write_var_int64(blen)
+    self.count += blen
+
+  cpdef write_byte(self, unsigned char _):
+    self.count += 1
+
+  cpdef write_bigendian_int64(self, libc.stdint.int64_t _):
+    self.count += 8
+
+  cpdef write_bigendian_int32(self, libc.stdint.int32_t _):
+    self.count += 4
+
+  cpdef size_t get_count(self):
+    return self.count
+
+  cpdef bytes get(self):
+    raise NotImplementedError
+
+  def __str__(self):
+    return '<%s %s>' % (self.__class__.__name__, self.count)
+
+
 cdef class InputStream(object):
   """An input string stream implementation supporting read() and size()."""
 

google/cloud/dataflow/coders/stream_test.py

@@ -14,6 +14,7 @@
 
 """Tests for the stream implementations."""
 
+import logging
 import math
 import unittest
 
@@ -22,8 +23,11 @@
 
 
 class StreamTest(unittest.TestCase):
+  # pylint: disable=invalid-name
   InputStream = slow_stream.InputStream
   OutputStream = slow_stream.OutputStream
+  ByteCountingOutputStream = slow_stream.ByteCountingOutputStream
+  # pylint: enable=invalid-name
 
   def test_read_write(self):
     out_s = self.OutputStream()
@@ -99,6 +103,28 @@ def test_read_write_bigendian_int32(self):
     for v in values:
       self.assertEquals(v, in_s.read_bigendian_int32())
 
+  def test_byte_counting(self):
+    bc_s = self.ByteCountingOutputStream()
+    self.assertEquals(0, bc_s.get_count())
+    bc_s.write('def')
+    self.assertEquals(3, bc_s.get_count())
+    bc_s.write('')
+    self.assertEquals(3, bc_s.get_count())
+    bc_s.write_byte(10)
+    self.assertEquals(4, bc_s.get_count())
+    # "nested" also writes the length of the string, which should
+    # cause 1 extra byte to be counted.
+    bc_s.write('2345', nested=True)
+    self.assertEquals(9, bc_s.get_count())
+    bc_s.write_var_int64(63)
+    self.assertEquals(10, bc_s.get_count())
+    bc_s.write_bigendian_int64(42)
+    self.assertEquals(18, bc_s.get_count())
+    bc_s.write_bigendian_int32(36)
+    self.assertEquals(22, bc_s.get_count())
+    bc_s.write_bigendian_double(6.25)
+    self.assertEquals(30, bc_s.get_count())
+
 
 try:
   # pylint: disable=g-import-not-at-top
@@ -108,22 +134,26 @@ class FastStreamTest(StreamTest):
     """Runs the test with the compiled stream classes."""
     InputStream = stream.InputStream
     OutputStream = stream.OutputStream
+    ByteCountingOutputStream = stream.ByteCountingOutputStream
 
 
   class SlowFastStreamTest(StreamTest):
     """Runs the test with compiled and uncompiled stream classes."""
     InputStream = stream.InputStream
     OutputStream = slow_stream.OutputStream
+    ByteCountingOutputStream = slow_stream.ByteCountingOutputStream
 
 
   class FastSlowStreamTest(StreamTest):
     """Runs the test with uncompiled and compiled stream classes."""
     InputStream = slow_stream.InputStream
     OutputStream = stream.OutputStream
+    ByteCountingOutputStream = stream.ByteCountingOutputStream
 
 except ImportError:
   pass
 
 
 if __name__ == '__main__':
+  logging.getLogger().setLevel(logging.INFO)
   unittest.main()