From a2ef82eac3c76210eb5f1390d7a0f0be9d9e2cc8 Mon Sep 17 00:00:00 2001
From: iguinn <iguinn@email.unc.edu>
Date: Tue, 29 Oct 2024 08:39:32 -0700
Subject: [PATCH] Added types to read_size_in_bytes; clarified some docstrings

---
 src/lgdo/lh5/_serializers/read/utils.py | 13 +++++++++----
 src/lgdo/lh5/store.py                   |  4 ++--
 src/lgdo/lh5/utils.py                   |  2 +-
 3 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/lgdo/lh5/_serializers/read/utils.py b/src/lgdo/lh5/_serializers/read/utils.py
index 3809483..3e2e8b9 100644
--- a/src/lgdo/lh5/_serializers/read/utils.py
+++ b/src/lgdo/lh5/_serializers/read/utils.py
@@ -104,7 +104,7 @@ def read_n_rows(h5o, fname, oname):
 
 
 def read_size_in_bytes(h5o, fname, oname, field_mask=None):
-    """Read number of rows in LH5 object"""
+    """Read number size in LH5 object in memory (in B)"""
     if not h5py.h5a.exists(h5o, b"datatype"):
         msg = "missing 'datatype' attribute"
         raise LH5DecodeError(msg, fname, oname)
@@ -116,11 +116,16 @@ def read_size_in_bytes(h5o, fname, oname, field_mask=None):
     lgdotype = datatype.datatype(type_attr)
 
     # scalars are dim-0 datasets
-    if lgdotype in (types.Scalar, types.Array, types.ArrayOfEqualSizedArrays):
+    if lgdotype in (
+        types.Scalar,
+        types.Array,
+        types.ArrayOfEqualSizedArrays,
+        types.FixedSizeArray,
+    ):
         return int(np.prod(h5o.shape) * h5o.dtype.itemsize)
 
     # structs don't have rows
-    if lgdotype is types.Struct:
+    if lgdotype in (types.Struct, types.Histogram, types.Histogram.Axis):
         size = 0
         for key in h5o:
             obj = h5py.h5o.open(h5o, key)
@@ -129,7 +134,7 @@ def read_size_in_bytes(h5o, fname, oname, field_mask=None):
         return size
 
     # tables should have elements with all the same length
-    if lgdotype is types.Table:
+    if lgdotype in (types.Table, types.WaveformTable):
         # read out each of the fields
         size = 0
         if not field_mask:
diff --git a/src/lgdo/lh5/store.py b/src/lgdo/lh5/store.py
index c8d876d..d3f488f 100644
--- a/src/lgdo/lh5/store.py
+++ b/src/lgdo/lh5/store.py
@@ -314,7 +314,7 @@ def read_n_rows(self, name: str, lh5_file: str | h5py.File) -> int | None:
         return utils.read_n_rows(name, self.gimme_file(lh5_file, "r"))
 
     def read_size_in_bytes(self, name: str, lh5_file: str | h5py.File) -> int:
-        """Look up the size (in B) of the object. Will recursively crawl
-        through all objects in a Struct or Table
+        """Look up the size (in B) of the object in memory. Will recursively
+        crawl through all objects in a Struct or Table
         """
         return utils.read_size_in_bytes(name, self.gimme_file(lh5_file, "r"))
diff --git a/src/lgdo/lh5/utils.py b/src/lgdo/lh5/utils.py
index 9ab0885..d86f0a3 100644
--- a/src/lgdo/lh5/utils.py
+++ b/src/lgdo/lh5/utils.py
@@ -56,7 +56,7 @@ def read_n_rows(name: str, h5f: str | h5py.File) -> int | None:
 
 
 def read_size_in_bytes(name: str, h5f: str | h5py.File) -> int | None:
-    """Look up the size (in B) in an LGDO object on disk. Will crawl
+    """Look up the size (in B) in an LGDO object in memory. Will crawl
     recursively through members of a Struct or Table
     """
     if not isinstance(h5f, h5py.File):