From ca6441caba6c6c0c8a0ae3e3f8eea01b67c31d76 Mon Sep 17 00:00:00 2001
From: Martin Pecka <peckama2@fel.cvut.cz>
Date: Fri, 25 Nov 2022 19:04:36 +0100
Subject: [PATCH] Better support for serializing buffer protocol objects. Added
 tests.

Signed-off-by: Martin Pecka <peckama2@fel.cvut.cz>
---
 src/genpy/generate_struct.py             |  14 +++
 src/genpy/generator.py                   |  26 +++--
 src/genpy/message.py                     |  16 ++-
 test/files/array/uint8_fixed_ser.txt     |  21 +++-
 test/files/array/uint8_fixed_ser_np.txt  |  21 +++-
 test/files/array/uint8_varlen_ser.txt    |  17 ++-
 test/files/array/uint8_varlen_ser_np.txt |  17 ++-
 test/msg/TestBinary.msg                  |   2 +
 test/test_genpy_generator.py             |  17 ++-
 test/test_genpy_message.py               | 135 +++++++++++++++++++++++
 10 files changed, 253 insertions(+), 33 deletions(-)
 create mode 100644 test/msg/TestBinary.msg

diff --git a/src/genpy/generate_struct.py b/src/genpy/generate_struct.py
index e6fa8f3..eac107a 100644
--- a/src/genpy/generate_struct.py
+++ b/src/genpy/generate_struct.py
@@ -186,3 +186,17 @@ def unpack3(var, struct_var, buff):
     :param buff: buffer that the unpack reads from, ``StringIO``
     """
     return '%s = %s.unpack(%s)' % (var, struct_var, buff)
+
+
+def memoryview_len(view):
+    """
+    Compute the size (in bytes) of a ``memoryview`` object.
+
+    This is the same as memoryview.nbytes, but compatible with Python 2.
+
+    :param view: The ``memoryview`` object.
+    """
+    length = view.itemsize
+    for s in view.shape:
+        length *= s
+    return length
diff --git a/src/genpy/generator.py b/src/genpy/generator.py
index e972aa9..f5c56ef 100644
--- a/src/genpy/generator.py
+++ b/src/genpy/generator.py
@@ -470,27 +470,37 @@ def string_serializer_generator(package, type_, name, serialize):  # noqa: D401
         # check to see if its a uint8/byte type, in which case we need to convert to string before serializing
         base_type, is_array, array_len = genmsg.msgs.parse_type(type_)
         if base_type in ['uint8', 'char']:
-            yield '# - if encoded as a list instead, serialize as bytes instead of string'
+            yield '# check for buffer protocol support'
             if array_len is None:
                 yield 'try:'
-                yield INDENT+'tmp = memoryview(%s) if python3 else buffer(%s)' % (var, var)
-                yield INDENT+pack2("'<I'", "tmp.nbytes if python3 else len(tmp)")
+                yield INDENT+'tmp = memoryview(%s)' % var
+                yield INDENT+'from genpy.generate_struct import memoryview_len'
+                yield INDENT+pack2("'<I'", "memoryview_len(tmp)")
                 yield INDENT+serializeit('tmp')
-                yield 'except Exception:'
+                yield 'except TypeError:'
+                yield INDENT+'# - if encoded as a list instead, serialize as bytes instead of string'
                 yield INDENT+'if type(%s) in [list, tuple]:' % var
                 yield INDENT+INDENT+pack2("'<I%sB'%length", 'length, *%s' % var)
                 yield INDENT+'else:'
                 yield INDENT+INDENT+pack2("'<I%ss'%length", 'length, %s' % var)
             else:
                 yield 'try:'
-                yield INDENT+'tmp = memoryview(%s)[:%s] if python3 else buffer(%s, 0, %s)' % (
-                    var, array_len, var, array_len)
-                yield INDENT+serializeit('tmp')
-                yield 'except Exception:'
+                yield INDENT+'# we want to process str and bytes by the except clause'
+                yield INDENT+'if isinstance(%s, bytes) or isinstance(%s, str):' % (var, var)
+                yield INDENT+INDENT+'raise TypeError()'
+                yield INDENT+'tmp = memoryview(%s)' % var
+                yield 'except TypeError:'
+                yield INDENT+'# - if encoded as a list instead, serialize as bytes instead of string'
                 yield INDENT+'if type(%s) in [list, tuple]:' % var
                 yield INDENT+INDENT+pack('%sB' % array_len, '*%s' % var)
                 yield INDENT+'else:'
                 yield INDENT+INDENT+pack('%ss' % array_len, var)
+                yield 'else:'
+                yield INDENT+'from genpy.generate_struct import memoryview_len'
+                yield INDENT+'tmp_len = memoryview_len(tmp)'
+                yield INDENT+'if tmp_len != %i:' % array_len
+                yield INDENT+INDENT+'raise TypeError("expected %i bytes, got %%i" %% (tmp_len,))' % array_len
+                yield INDENT+serializeit('tmp[:%i]' % array_len)
         else:
             # FIXME: for py3k, this needs to be w/ encode(), but this interferes with actual byte data
             yield 'if python3 or type(%s) == unicode:' % (var)
diff --git a/src/genpy/message.py b/src/genpy/message.py
index d61ab0f..26ac074 100644
--- a/src/genpy/message.py
+++ b/src/genpy/message.py
@@ -48,6 +48,7 @@
 import yaml
 
 from .base import is_simple
+from .generate_struct import memoryview_len
 from .rostime import Duration
 from .rostime import TVal
 from .rostime import Time
@@ -299,14 +300,25 @@ def check_type(field_name, field_type, field_val):
         base_type = field_type[:field_type.index('[')]
 
         if type(field_val) in (bytes, str):
-            if base_type not in ['char', 'uint8']:
+            if base_type not in ('char', 'uint8'):
                 raise SerializationError('field %s must be a list or tuple type. Only uint8[] can be a string' % field_name)
             else:
                 # It's a string so its already in byte format and we
                 # don't need to check the individual bytes in the
                 # string.
                 return
-
+        if base_type in ('char', 'uint8'):
+            # check buffer protocol support
+            try:
+                length = memoryview_len(memoryview(field_val))
+                _, _, array_len = genmsg.msgs.parse_type(field_type)
+                if array_len is not None and length != array_len:
+                    raise SerializationError('field %s must receive %i bytes, but %i given' % (
+                        field_name, array_len, length))
+                # if the value supports buffer protocol, pass it directly
+                return
+            except TypeError:
+                pass  # does not support buffer protocol
         if not type(field_val) in [list, tuple]:
             raise SerializationError('field %s must be a list or tuple type' % field_name)
         for v in field_val:
diff --git a/test/files/array/uint8_fixed_ser.txt b/test/files/array/uint8_fixed_ser.txt
index 803d42b..893e23a 100644
--- a/test/files/array/uint8_fixed_ser.txt
+++ b/test/files/array/uint8_fixed_ser.txt
@@ -1,5 +1,18 @@
-# - if encoded as a list instead, serialize as bytes instead of string
-if type(data) in [list, tuple]:
-  buff.write(_get_struct_8B().pack(*data))
+# check for buffer protocol support
+try:
+  # we want to process str and bytes by the except clause
+  if isinstance(data, bytes) or isinstance(data, str):
+    raise TypeError()
+  tmp = memoryview(data)
+except TypeError:
+  # - if encoded as a list instead, serialize as bytes instead of string
+  if type(data) in [list, tuple]:
+    buff.write(_get_struct_8B().pack(*data))
+  else:
+    buff.write(_get_struct_8s().pack(data))
 else:
-  buff.write(_get_struct_8s().pack(data))
+  from genpy.generate_struct import memoryview_len
+  tmp_len = memoryview_len(tmp)
+  if tmp_len != 8:
+    raise TypeError("expected 8 bytes, got %i" % (tmp_len,))
+  buff.write(tmp[:8])
\ No newline at end of file
diff --git a/test/files/array/uint8_fixed_ser_np.txt b/test/files/array/uint8_fixed_ser_np.txt
index 803d42b..893e23a 100644
--- a/test/files/array/uint8_fixed_ser_np.txt
+++ b/test/files/array/uint8_fixed_ser_np.txt
@@ -1,5 +1,18 @@
-# - if encoded as a list instead, serialize as bytes instead of string
-if type(data) in [list, tuple]:
-  buff.write(_get_struct_8B().pack(*data))
+# check for buffer protocol support
+try:
+  # we want to process str and bytes by the except clause
+  if isinstance(data, bytes) or isinstance(data, str):
+    raise TypeError()
+  tmp = memoryview(data)
+except TypeError:
+  # - if encoded as a list instead, serialize as bytes instead of string
+  if type(data) in [list, tuple]:
+    buff.write(_get_struct_8B().pack(*data))
+  else:
+    buff.write(_get_struct_8s().pack(data))
 else:
-  buff.write(_get_struct_8s().pack(data))
+  from genpy.generate_struct import memoryview_len
+  tmp_len = memoryview_len(tmp)
+  if tmp_len != 8:
+    raise TypeError("expected 8 bytes, got %i" % (tmp_len,))
+  buff.write(tmp[:8])
\ No newline at end of file
diff --git a/test/files/array/uint8_varlen_ser.txt b/test/files/array/uint8_varlen_ser.txt
index 2e5f8e8..a68a4bf 100644
--- a/test/files/array/uint8_varlen_ser.txt
+++ b/test/files/array/uint8_varlen_ser.txt
@@ -1,6 +1,13 @@
 length = len(data)
-# - if encoded as a list instead, serialize as bytes instead of string
-if type(data) in [list, tuple]:
-  buff.write(struct.Struct('<I%sB'%length).pack(length, *data))
-else:
-  buff.write(struct.Struct('<I%ss'%length).pack(length, data))
+# check for buffer protocol support
+try:
+  tmp = memoryview(data)
+  from genpy.generate_struct import memoryview_len
+  buff.write(struct.Struct('<I').pack(memoryview_len(tmp)))
+  buff.write(tmp)
+except TypeError:
+  # - if encoded as a list instead, serialize as bytes instead of string
+  if type(data) in [list, tuple]:
+    buff.write(struct.Struct('<I%sB'%length).pack(length, *data))
+  else:
+    buff.write(struct.Struct('<I%ss'%length).pack(length, data))
\ No newline at end of file
diff --git a/test/files/array/uint8_varlen_ser_np.txt b/test/files/array/uint8_varlen_ser_np.txt
index 2e5f8e8..a68a4bf 100644
--- a/test/files/array/uint8_varlen_ser_np.txt
+++ b/test/files/array/uint8_varlen_ser_np.txt
@@ -1,6 +1,13 @@
 length = len(data)
-# - if encoded as a list instead, serialize as bytes instead of string
-if type(data) in [list, tuple]:
-  buff.write(struct.Struct('<I%sB'%length).pack(length, *data))
-else:
-  buff.write(struct.Struct('<I%ss'%length).pack(length, data))
+# check for buffer protocol support
+try:
+  tmp = memoryview(data)
+  from genpy.generate_struct import memoryview_len
+  buff.write(struct.Struct('<I').pack(memoryview_len(tmp)))
+  buff.write(tmp)
+except TypeError:
+  # - if encoded as a list instead, serialize as bytes instead of string
+  if type(data) in [list, tuple]:
+    buff.write(struct.Struct('<I%sB'%length).pack(length, *data))
+  else:
+    buff.write(struct.Struct('<I%ss'%length).pack(length, data))
\ No newline at end of file
diff --git a/test/msg/TestBinary.msg b/test/msg/TestBinary.msg
new file mode 100644
index 0000000..84d3cf4
--- /dev/null
+++ b/test/msg/TestBinary.msg
@@ -0,0 +1,2 @@
+uint8[] data_var
+uint8[4] data_fixed
diff --git a/test/test_genpy_generator.py b/test/test_genpy_generator.py
index a7b2493..59431fc 100644
--- a/test/test_genpy_generator.py
+++ b/test/test_genpy_generator.py
@@ -318,11 +318,18 @@ def test_string_serializer_generator():
     for t in ['uint8[]', 'byte[]', 'uint8[10]', 'byte[20]']:
         g = genpy.generator.string_serializer_generator('foo', 'uint8[]', 'b_name', True)
         assert """length = len(b_name)
-# - if encoded as a list instead, serialize as bytes instead of string
-if type(b_name) in [list, tuple]:
-  buff.write(struct.Struct('<I%sB'%length).pack(length, *b_name))
-else:
-  buff.write(struct.Struct('<I%ss'%length).pack(length, b_name))""" == '\n'.join(g)
+# check for buffer protocol support
+try:
+  tmp = memoryview(b_name)
+  from genpy.generate_struct import memoryview_len
+  buff.write(struct.Struct('<I').pack(memoryview_len(tmp)))
+  buff.write(tmp)
+except TypeError:
+  # - if encoded as a list instead, serialize as bytes instead of string
+  if type(b_name) in [list, tuple]:
+    buff.write(struct.Struct('<I%sB'%length).pack(length, *b_name))
+  else:
+    buff.write(struct.Struct('<I%ss'%length).pack(length, b_name))""" == '\n'.join(g)
 
     # Test Deserializers
     val = """start = end
diff --git a/test/test_genpy_message.py b/test/test_genpy_message.py
index e890229..2ee6eef 100644
--- a/test/test_genpy_message.py
+++ b/test/test_genpy_message.py
@@ -367,10 +367,19 @@ def test_check_types_valid(self):
         import numpy as np
         genpy.message.check_type('test', 'uint8[]', 'byteDataIsAStringInPy')
         genpy.message.check_type('test', 'char[]', 'byteDataIsAStringInPy')
+        genpy.message.check_type('test', 'uint8[]', b'byteDataIsAStringInPy')
         genpy.message.check_type('test', 'uint8[]', [3, 4, 5])
         genpy.message.check_type('test', 'uint8[]', (3, 4, 5))
         genpy.message.check_type('test', 'char[]', [3, 4, 5])
         genpy.message.check_type('test', 'int32[]', [3, 4, 5])
+        genpy.message.check_type('test', 'uint8[]', bytearray(b'bytes'))
+        genpy.message.check_type('test', 'uint8[]', np.array([3, 4, 5]))
+        # serializing string as bytes does not check that correct number of bytes were given
+        genpy.message.check_type('test', 'uint8[3]', "asd".encode())
+        genpy.message.check_type('test', 'uint8[3]', "asda".encode())
+        genpy.message.check_type('test', 'uint8[3]', "as".encode())
+        # serializing buffer-protocol objects does check the number of bytes
+        genpy.message.check_type('test', 'uint8[3]', np.array([3, 4, 5], dtype=np.uint8))
         genpy.message.check_type('test', 'int32', -5)
         genpy.message.check_type('test', 'int64', -5)
         genpy.message.check_type('test', 'int16', -5)
@@ -404,10 +413,13 @@ def test_check_types_valid(self):
 
     def test_check_types_invalid(self):
         from genpy import SerializationError
+        import numpy as np
         self.assertRaises(SerializationError, genpy.message.check_type,
                           'test', 'int32[]', 'someString')
         self.assertRaises(SerializationError, genpy.message.check_type,
                           'test', 'uint32[]', [3, -2, 4])
+        self.assertRaises(SerializationError, genpy.message.check_type,
+                          'test', 'uint8[4]', np.array([3, -2, 4], dtype=np.uint8))
         self.assertRaises(SerializationError, genpy.message.check_type,
                           'test', 'uint8', -2)
         self.assertRaises(SerializationError, genpy.message.check_type,
@@ -694,9 +706,12 @@ def test_check_type(self):
         # test to validate this.
         from genpy.message import check_type, SerializationError
         from genpy import Time, Duration
+        import numpy as np
         valids = [
             ('byte', 1), ('byte', -1),
             ('string', ''), ('string', 'a string of text'),
+            ('uint8[]', "buffer"), ('uint8[]', b"buffer"), ('uint8[]', bytearray(b"buffer")),
+            ('uint8[]', np.array([1, 2, 3, 4])),
             ('int32[]', []),
             ('int32[]', [1, 2, 3, 4]),
             ('time', Time()), ('time', Time.from_sec(1.0)),
@@ -720,6 +735,8 @@ def test_check_type(self):
             ('uint8', -1), ('uint8', 112312),
             ('int32', '1'), ('int32', 1.),
             ('int32[]', 1), ('int32[]', [1., 2.]), ('int32[]', [1, 2.]),
+            ('uint16[]', "buffer"), ('uint16[]', b"buffer"), ('uint16[]', bytearray(b"buffer")),
+            ('uint16[]', np.array([1, 2, 3, 4])),
             ('duration', 1), ('time', 1),
             ]
         for t, v in invalids:
@@ -757,6 +774,124 @@ def test_serialize_exception_msg(self):
         except Exception:
             assert False, 'This should have raised a genpy.SerializationError instead'
 
+    def test_serialize_binary_msg(self):
+        from genpy.msg import TestBinary
+        import numpy as np
+        try:
+            from cStringIO import StringIO
+        except ImportError:
+            from io import BytesIO as StringIO
+
+        m = TestBinary()
+
+        buff = StringIO()
+        m.data_var = [3, 4, 5]
+        m.data_fixed = [1, 2, 3, 4]
+        m.serialize(buff)
+        expected_val = buff.getvalue()
+
+        buff = StringIO()
+        m.data_var = (3, 4, 5)
+        m.data_fixed = (1, 2, 3, 4)
+        m.serialize(buff)
+        val = buff.getvalue()
+        self.assertEqual(val, expected_val)
+
+        buff = StringIO()
+        m.data_var = "\x03\x04\x05".encode()
+        m.data_fixed = "\x01\x02\x03\x04".encode()
+        m.serialize(buff)
+        val = buff.getvalue()
+        self.assertEqual(val, expected_val)
+
+        buff = StringIO()
+        m.data_var = b"\x03\x04\x05"
+        m.data_fixed = b"\x01\x02\x03\x04"
+        m.serialize(buff)
+        val = buff.getvalue()
+        self.assertEqual(val, expected_val)
+
+        buff = StringIO()
+        m.data_var = bytearray(b"\x03\x04\x05")
+        m.data_fixed = bytearray(b"\x01\x02\x03\x04")
+        m.serialize(buff)
+        val = buff.getvalue()
+        self.assertEqual(val, expected_val)
+
+        buff = StringIO()
+        m.data_var = np.array([3, 4, 5], dtype=np.uint8)
+        m.data_fixed = np.array([1, 2, 3, 4], dtype=np.uint8)
+        m.serialize(buff)
+        self.assertEqual(val, expected_val)
+
+        # structured numpy array with one 3- or 4-dimensional element (useful e.g. for pointclouds)
+        dtype_var = [('x', np.uint8), ('y', np.uint8), ('z', np.uint8)]
+        dtype_fixed = [('x', np.uint8), ('y', np.uint8), ('z', np.uint8), ('w', np.uint8)]
+
+        buff = StringIO()
+        m.data_var = np.zeros((1,), dtype=dtype_var)
+        m.data_var[0]['x'] = 3
+        m.data_var[0]['y'] = 4
+        m.data_var[0]['z'] = 5
+        m.data_fixed = np.zeros((1,), dtype=dtype_fixed)
+        m.data_fixed[0]['x'] = 1
+        m.data_fixed[0]['y'] = 2
+        m.data_fixed[0]['z'] = 3
+        m.data_fixed[0]['w'] = 4
+        m.serialize(buff)
+        self.assertEqual(val, expected_val)
+
+        # str and bytes iterables do not care about exact sizes for fixed fields
+
+        buff = StringIO()
+        m.data_var = "\x03\x04\x05".encode()
+        m.data_fixed = "\x01\x02\x03\x04\x05".encode()
+        m.serialize(buff)
+        expected_val = buff.getvalue()
+
+        buff = StringIO()
+        m.data_var = b"\x03\x04\x05"
+        m.data_fixed = b"\x01\x02\x03\x04\x05"
+        m.serialize(buff)
+        val = buff.getvalue()
+        self.assertEqual(val, expected_val)
+
+        # buffer-protocol values do care about exact size for fixed-size arrays
+
+        buff = StringIO()
+        m.data_var = bytearray(b"\x03\x04\x05")
+        m.data_fixed = bytearray(b"\x01\x02\x03\x04\x05")
+        try:
+            m.serialize(buff)
+            assert False, 'This should have raised a genpy.SerializationError'
+        except genpy.SerializationError as e:
+            self.assertEqual(str(e), "field data_fixed must receive 4 bytes, but 5 given")
+        except Exception:
+            assert False, 'This should have raised a genpy.SerializationError instead'
+
+        buff = StringIO()
+        m.data_var = np.array([3, 4, 5, 6], dtype=np.uint8)
+        m.data_fixed = np.array([1, 2, 3, 4, 5], dtype=np.uint8)
+        try:
+            m.serialize(buff)
+            assert False, 'This should have raised a genpy.SerializationError'
+        except genpy.SerializationError as e:
+            self.assertEqual(str(e), "field data_fixed must receive 4 bytes, but 5 given")
+        except Exception:
+            assert False, 'This should have raised a genpy.SerializationError instead'
+
+        # wrong dtype (default is np.float64) results in a different count of bytes
+        buff = StringIO()
+        m.data_var = np.array([3, 4, 5])
+        m.data_fixed = np.array([1, 2, 3, 4])
+        try:
+            m.serialize(buff)
+            assert False, 'This should have raised a genpy.SerializationError'
+        except genpy.SerializationError as e:
+            self.assertEqual(str(e), "field data_fixed must receive 4 bytes, but 32 given")
+        except Exception:
+            assert False, 'This should have raised a genpy.SerializationError instead'
+
     @unittest.skipIf(sys.hexversion < 0x03000000, "Python 3 only test")
     def test_deserialize_unicode_error(self):
         from genpy.msg import TestString, TestMsgArray