Test case for corner case file (buffer too small)

pyfive/dataobjects.py

@@ -178,6 +178,10 @@ def _get_attributes_from_attr_info(self, attrs, attr_info):
         adict = dict()
         for record in btree.iter_records():
             data = heap.get_data(record['heapid'])
+            # in the unlikely event of loss of data
+            # just move on
+            if data is None:
+                continue
             name, value = self._parse_attribute_msg(data,0)
             adict[name] = value
         return adict

pyfive/misc_low_level.py

@@ -229,23 +229,48 @@ def __init__(self, fh, offset):
         self.nobjects = header["managed_object_count"] + header["huge_object_count"] + header["tiny_object_count"]
 
         managed = []
+        # while iterating over direct and indirect blocks we keep track of the heap_offset
+        # thus, we are able to map this later back to an offset into our managed heap buffer
+        blocks = []
+        buffer_offset = 0
         root_address = header["root_block_address"]
         if root_address:
             nrows = header["indirect_current_rows_count"]
             if nrows:
-                for data in self._iter_indirect_block(fh, root_address, nrows):
+                for data, heap_offset, block_size in self._iter_indirect_block(fh, root_address, nrows):
                     managed.append(data)
+                    blocks.append((heap_offset, buffer_offset, block_size))
+                    buffer_offset += len(data)
             else:
-                data = self._read_direct_block(fh, root_address, start_block_size)
+                data, heap_offset = self._read_direct_block(fh, root_address, start_block_size)
                 managed.append(data)
+                blocks.append((heap_offset, buffer_offset, start_block_size))
+                buffer_offset += len(data)
+
         self.managed = b"".join(managed)
+        self.blocks = blocks
 
     def _read_direct_block(self, fh, offset, block_size):
+        """
+        Read FHDB - direct block - from heap and return data and heap offset
+        """
         fh.seek(offset)
         data = fh.read(block_size)
         header = _unpack_struct_from(self.direct_block_header, data)
-        header["signature"] == b"FHDB"
-        return data
+        assert header["signature"] == b"FHDB"
+        return data, int.from_bytes(header["block_offset"],
+                                                byteorder="little", signed=False)
+
+    def heapid_to_buffer_offset(self, heapid_offset):
+        """
+        Get offset into flat managed buffer from heapid offset
+        """
+        for heap_offset, buffer_offset, block_size in self.blocks:
+            if heap_offset <= heapid_offset < heap_offset + block_size:
+                relative = heapid_offset - heap_offset
+                return buffer_offset + relative
+
+        raise KeyError("HeapID offset not inside any heap block")
 
     def get_data(self, heapid):
         firstbyte = heapid[0]
@@ -262,7 +287,14 @@ def get_data(self, heapid):
                 data_offset += nbytes
                 nbytes = self._managed_object_length_size
                 size = _unpack_integer(nbytes, heapid, data_offset)
-                return self.managed[offset:offset+size]
+
+                # map heap_id offset to flat buffer offset
+                offset = self.heapid_to_buffer_offset(offset)
+                if offset < len(self.managed):
+                    return self.managed[offset:offset + size]
+
+                return None
+
             case 1: # tiny
                 raise NotImplementedError
             case 2: # huge
@@ -296,26 +328,26 @@ def _iter_indirect_block(self, fh, offset, nrows):
         for i in range(ndirect):
             address = struct.unpack('<Q', fh.read(8))[0]
             if address == UNDEFINED_ADDRESS:
-                break
+                continue
             block_size = self._calc_block_size(i)
             direct_blocks.append((address, block_size))
 
         indirect_blocks = list()
         for i in range(ndirect, ndirect+nindirect):
             address = struct.unpack('<Q', fh.read(8))[0]
             if address == UNDEFINED_ADDRESS:
-                break
+                continue
             block_size = self._calc_block_size(i)
             nrows = self._iblock_nrows_from_block_size(block_size)
-            indirect_blocks.append((address, nrows))
+            indirect_blocks.append((address, block_size, nrows))
 
         for address, block_size in direct_blocks:
-            obj = self._read_direct_block(fh, address, block_size)
-            yield obj
+            obj, heap_offset = self._read_direct_block(fh, address, block_size)
+            yield obj, heap_offset, block_size
 
-        for address, nrows in indirect_blocks:
-            for obj in self._iter_indirect_block(fh, address, nrows):
-                yield obj
+        for address, block_size, nrows in indirect_blocks:
+            for obj, heap_offset in self._iter_indirect_block(fh, address, nrows):
+                yield obj, heap_offset, block_size
 
     def _calc_block_size(self, iblock):
         row = iblock//self.header["table_width"]

tests/test_buffer_issue.py

@@ -36,3 +36,14 @@ def test_buffer_issue():
     print("File with issue da193a_25_6hr_t_pt_cordex__198807-198807.nc")
     print("Variable m01s30i111")
     _load_nc_file('m01s30i111')
+
+
+def test_buffer_issue_ukesm():
+    """Test with yet another corner case file."""
+    fp = "tests/data/noy_AERmonZ_UKESM1-0-LL_piControl_r1i1p1f2_gnz_200001-200012.nc"
+    with pyfive.File(fp) as pfile:
+        print(pfile["noy"])
+        attrs = pfile["noy"].attrs
+        print(len(attrs))
+        print(attrs.keys())
+