From d943ee80b9968dac9bfab031676d976e7d4d1963 Mon Sep 17 00:00:00 2001
From: Ruben Verweij <verweij@physics.leidenuniv.nl>
Date: Tue, 18 Jul 2017 17:00:15 +0200
Subject: [PATCH] Unit testing: test the correct parsing of image_attributes

---
 nd2reader/artificial.py    | 88 ++++++++++++++++++++++++++++++++++++--
 nd2reader/common.py        |  7 +--
 tests/test_raw_metadata.py | 15 +++++++
 3 files changed, 101 insertions(+), 9 deletions(-)

diff --git a/nd2reader/artificial.py b/nd2reader/artificial.py
index 3a0a8f9..ac7ec8c 100644
--- a/nd2reader/artificial.py
+++ b/nd2reader/artificial.py
@@ -12,6 +12,15 @@ class ArtificialND2(object):
     """
     header = 0xabeceda
     relative_offset = 0
+    data_types = {'unsigned_char': 1,
+                  'unsigned_int': 2,
+                  'unsigned_int_2': 3,
+                  'unsigned_long': 5,
+                  'double': 6,
+                  'string': 8,
+                  'char_array': 9,
+                  'metadata_item': 11,
+                  }
 
     def __init__(self, file, version=(3, 0)):
         self.version = version
@@ -151,14 +160,85 @@ class ArtificialND2(object):
         return raw_text, locations, file_data_dict
 
     def _pack_data_with_metadata(self, data):
-        data = struct.pack('I', data)
-        raw_data = struct.pack("IIQ", self.header, self.relative_offset, len(data))
-        raw_data += data
+        packed_data = self._pack_raw_data_with_metadata(data)
+
+        raw_data = struct.pack("IIQ", self.header, self.relative_offset, len(packed_data))
+        raw_data += packed_data
+
+        return raw_data
+
+    def _pack_raw_data_with_metadata(self, data):
+        raw_data = b''
+
+        if isinstance(data, dict):
+            raw_data = self._pack_dict_with_metadata(data)
+        elif isinstance(data, int):
+            raw_data = struct.pack('I', data)
+        elif isinstance(data, float):
+            raw_data = struct.pack('d', data)
+        elif isinstance(data, str):
+            raw_data = bytes(data, 'utf-8')
+
+        return raw_data
+
+    def _get_data_type(self, data):
+        if isinstance(data, dict):
+            return self.data_types['metadata_item']
+        elif isinstance(data, int):
+            return self.data_types['unsigned_int']
+        elif isinstance(data, float):
+            return self.data_types['double']
+        elif isinstance(data, str):
+            return self.data_types['string']
+
+        return np.max(self.data_types.values()) + 1
+
+    def _pack_dict_with_metadata(self, data):
+        raw_data = b''
+
+        for data_key in data.keys():
+            # names have always one character extra and are padded in zero bytes???
+            b_data_key = b''.join([struct.pack('cx', bytes(s, 'utf-8')) for s in data_key]) + struct.pack('xx')
+
+            # header consists of data type and length of key name, it is represented by 2 unsigned chars
+            raw_data += struct.pack('BB', self._get_data_type(data[data_key]), len(data_key) + 1)
+            raw_data += b_data_key
+
+            sub_data = self._pack_raw_data_with_metadata(data[data_key])
+
+            if isinstance(data[data_key], dict):
+                # Pack: the number of keys and the length of raw data until now, sub data
+                # and the 12 bytes that we add now
+                raw_data += struct.pack("<IQ", len(data[data_key].keys()), len(sub_data) + len(raw_data) + 12)
+
+            raw_data += sub_data
+
+            if isinstance(data[data_key], dict):
+                # apparently there is also a huge empty space
+                raw_data += b''.join([struct.pack('x')] * len(data[data_key].keys()) * 8)
+
         return raw_data
 
     def _get_file_data(self, labels):
         file_data = [
-            7,  # ImageAttributesLV!",
+            {
+                'SLxImageAttributes':
+                    {
+                        'uiWidth': 128,
+                        'uiWidthBytes': 256,
+                        'uiHeight': 128,
+                        'uiComp': 1,
+                        'uiBpcInMemory': 16,
+                        'uiBpcSignificant': 12,
+                        'uiSequenceCount': 70,
+                        'uiTileWidth': 128,
+                        'uiTileHeight': 128,
+                        'eCompression': 2,
+                        'dCompressionParam': -1.0,
+                        'ePixelType': 1,
+                        'uiVirtualComponents': 1
+                    }
+            },  # ImageAttributesLV!",
             7,  # ImageTextInfoLV!",
             7,  # ImageMetadataLV!",
             7,  # ImageMetadataSeqLV|0!",
diff --git a/nd2reader/common.py b/nd2reader/common.py
index 52fba85..e9842d5 100644
--- a/nd2reader/common.py
+++ b/nd2reader/common.py
@@ -257,7 +257,7 @@ def _get_value(data, data_type, cursor_position):
               11: _parse_metadata_item}
     try:
         value = parser[data_type](data) if data_type < 11 else parser[data_type](data, cursor_position)
-    except KeyError:
+    except (KeyError, struct.error):
         value = None
 
     return value
@@ -289,10 +289,7 @@ def read_metadata(data, count):
             # We've reached the end of some hierarchy of data
             break
 
-        if six.PY3:
-            header = header.decode("utf8")
-
-        data_type, name_length = map(ord, header)
+        data_type, name_length = struct.unpack('BB', header)
         name = data.read(name_length * 2).decode("utf16")[:-1].encode("utf8")
         value = _get_value(data, data_type, cursor_position)
 
diff --git a/tests/test_raw_metadata.py b/tests/test_raw_metadata.py
index 0269ecb..8703980 100644
--- a/tests/test_raw_metadata.py
+++ b/tests/test_raw_metadata.py
@@ -36,3 +36,18 @@ class TestRawMetadata(unittest.TestCase):
 
     def test_pfs_status(self):
         self.assertEqual(self.file_data['pfs_status'], self.metadata.pfs_status[0])
+
+    def _assert_dicts_equal(self, parsed_dict, original_dict):
+        for attribute in original_dict.keys():
+            parsed_key = bytes(attribute, 'utf-8')
+            self.assertIn(parsed_key, parsed_dict.keys())
+
+            if isinstance(parsed_dict[parsed_key], dict):
+                self._assert_dicts_equal(parsed_dict[parsed_key], original_dict[attribute])
+            else:
+                self.assertEqual(parsed_dict[parsed_key], original_dict[attribute])
+
+    def test_image_attributes(self):
+        parsed_dict = self.metadata.image_attributes
+
+        self._assert_dicts_equal(parsed_dict, self.file_data['image_attributes'])