[compression] DEFAULT_HDF5_COMPRESSION can be a dict of options

legend-exp · Nov 8, 2023 · b41bac4 · b41bac4
1 parent 77b4b43
commit b41bac4
Showing 1 changed file with 12 additions and 8 deletions.
diff --git a/src/lgdo/lh5_store.py b/src/lgdo/lh5_store.py
@@ -38,7 +38,7 @@
 
 log = logging.getLogger(__name__)
 
-DEFAULT_HDF5_COMPRESSION = None
+DEFAULT_HDF5_COMPRESSION = "gzip"
 
 
 class LH5Store:
@@ -1091,18 +1091,22 @@ def write_object(
                     log.debug(f"overwriting {name} in {group}")
                     del group[name]
 
-                # create HDF5 dataset
+                # set default compression options
+                if isinstance(DEFAULT_HDF5_COMPRESSION, dict):
+                    for k, v in DEFAULT_HDF5_COMPRESSION.items():
+                        h5py_kwargs.setdefault(k, v)
+                else:
+                    h5py_kwargs.setdefault("compression", DEFAULT_HDF5_COMPRESSION)
+
                 # compress using the 'compression' LGDO attribute, if available
                 if "compression" in obj.attrs:
                     comp_algo = obj.attrs["compression"]
-                    if not isinstance(comp_algo, dict):
-                        h5py_kwargs["compression"] = comp_algo
+                    if isinstance(comp_algo, dict):
+                        h5py_kwargs |= obj.attrs["compression"]
                     else:
-                        h5py_kwargs |= comp_algo
-
-                # otherwise use "compression" argument
-                h5py_kwargs.setdefault("compression", DEFAULT_HDF5_COMPRESSION)
+                        h5py_kwargs["compression"] = obj.attrs["compression"]
 
+                # create HDF5 dataset
                 ds = group.create_dataset(
                     name, data=nda, maxshape=maxshape, **h5py_kwargs
                 )