legend-pydataobj 1.5.0a5__py3-none-any.whl → 1.6.0__py3-none-any.whl

lgdo/lh5/tools.py

@@ -2,11 +2,9 @@ from __future__ import annotations
 
 import fnmatch
 import glob
-import inspect
 import logging
 import os
-from collections.abc import Iterable
-from typing import Any
+from copy import copy
 from warnings import warn
 
 import h5py
@@ -20,7 +18,11 @@ from .store import LH5Store
 log = logging.getLogger(__name__)
 
 
-def ls(lh5_file: str | h5py.Group, lh5_group: str = "") -> list[str]:
+def ls(
+    lh5_file: str | h5py.Group,
+    lh5_group: str = "",
+    recursive: bool = False,
+) -> list[str]:
     """Return a list of LH5 groups in the input file and group, similar
     to ``ls`` or ``h5ls``. Supports wildcards in group names.
 
@@ -32,6 +34,8 @@ def ls(lh5_file: str | h5py.Group, lh5_group: str = "") -> list[str]:
     lh5_group
         group to search. add a ``/`` to the end of the group name if you want to
         list all objects inside that group.
+    recursive
+        if ``True``, recurse into subgroups.
     """
 
     log.debug(
@@ -49,15 +53,30 @@ def ls(lh5_file: str | h5py.Group, lh5_group: str = "") -> list[str]:
     if lh5_group == "":
         lh5_group = "*"
 
+    # get the first group in the group path
     splitpath = lh5_group.split("/", 1)
+    # filter out objects that don't match lh5_group pattern
     matchingkeys = fnmatch.filter(lh5_file.keys(), splitpath[0])
 
+    ret = []
+    # if there were no "/" in lh5_group just return the result
     if len(splitpath) == 1:
-        return matchingkeys
+        ret = matchingkeys
+
+    else:
+        for key in matchingkeys:
+            ret.extend([f"{key}/{path}" for path in ls(lh5_file[key], splitpath[1])])
+
+    if recursive:
+        rec_ret = copy(ret)
+        for obj in ret:
+            try:
+                rec_ret += ls(lh5_file, lh5_group=f"{obj}/", recursive=True)
+            except AttributeError:
+                continue
+
+        return rec_ret
 
-    ret = []
-    for key in matchingkeys:
-        ret.extend([f"{key}/{path}" for path in ls(lh5_file[key], splitpath[1])])
     return ret
 
 
@@ -169,51 +188,6 @@ def show(
         key = k_new
 
 
-def read_as(
-    name: str,
-    lh5_file: str | h5py.File | Iterable[str | h5py.File],
-    library: str,
-    **kwargs,
-) -> Any:
-    """Read LH5 data from disk straight into a third-party data format view.
-
-    This function is nothing more than a shortcut chained call to
-    :meth:`.LH5Store.read` and to :meth:`.LGDO.view_as`.
-
-    Parameters
-    ----------
-    name
-        LH5 object name on disk.
-    lh5_file
-        LH5 file name.
-    library
-        string ID of the third-party data format library (``np``, ``pd``,
-        ``ak``, etc).
-
-    See Also
-    --------
-    .LH5Store.read, .LGDO.view_as
-    """
-    # determine which keyword arguments should be forwarded to read() and which
-    # should be forwarded to view_as()
-    read_kwargs = inspect.signature(LH5Store.read).parameters.keys()
-
-    kwargs1 = {}
-    kwargs2 = {}
-    for k, v in kwargs.items():
-        if k in read_kwargs:
-            kwargs1[k] = v
-        else:
-            kwargs2[k] = v
-
-    # read the LGDO from disk
-    store = LH5Store()
-    obj, _ = store.read(name, lh5_file, **kwargs1)
-
-    # and finally return a view
-    return obj.view_as(library, **kwargs2)
-
-
 def load_nda(
     f_list: str | list[str],
     par_list: list[str],

lgdo/lh5/utils.py

@@ -1,46 +1,149 @@
 """Implements utilities for LEGEND Data Objects."""
+
 from __future__ import annotations
 
 import glob
 import logging
 import os
 import string
+from collections.abc import Mapping, Sequence
+from typing import Any
 
-log = logging.getLogger(__name__)
+import h5py
 
+from .. import types
+from . import _serializers, datatype
+from .exceptions import LH5DecodeError
 
-def parse_datatype(datatype: str) -> tuple[str, tuple[int, ...], str | list[str]]:
-    """Parse datatype string and return type, dimensions and elements.
+log = logging.getLogger(__name__)
 
-    Parameters
-    ----------
-    datatype
-        a LGDO-formatted datatype string.
 
-    Returns
-    -------
-    element_type
-        the datatype name dims if not ``None``, a tuple of dimensions for the
-        LGDO. Note this is not the same as the NumPy shape of the underlying
-        data object. See the LGDO specification for more information. Also see
-        :class:`~.types.ArrayOfEqualSizedArrays` and
-        :meth:`.lh5_store.LH5Store.read` for example code elements for
-        numeric objects, the element type for struct-like  objects, the list of
-        fields in the struct.
+def get_buffer(
+    name: str,
+    lh5_file: str | h5py.File | Sequence[str | h5py.File],
+    size: int | None = None,
+    field_mask: Mapping[str, bool] | Sequence[str] | None = None,
+) -> types.LGDO:
+    """Returns an LGDO appropriate for use as a pre-allocated buffer.
+
+    Sets size to `size` if object has a size.
     """
-    if "{" not in datatype:
-        return "scalar", None, datatype
+    obj, n_rows = _serializers._h5_read_lgdo(
+        name, lh5_file, n_rows=0, field_mask=field_mask
+    )
 
-    # for other datatypes, need to parse the datatype string
-    from parse import parse
+    if hasattr(obj, "resize") and size is not None:
+        obj.resize(new_size=size)
 
-    datatype, element_description = parse("{}{{{}}}", datatype)
-    if datatype.endswith(">"):
-        datatype, dims = parse("{}<{}>", datatype)
-        dims = [int(i) for i in dims.split(",")]
-        return datatype, tuple(dims), element_description
+    return obj
 
-    return datatype, None, element_description.split(",")
+
+def read_n_rows(name: str, h5f: str | h5py.File) -> int | None:
+    """Look up the number of rows in an Array-like LGDO object on disk.
+
+    Return ``None`` if `name` is a :class:`.Scalar` or a :class:`.Struct`.
+    """
+    if not isinstance(h5f, h5py.File):
+        h5f = h5py.File(h5f, "r")
+
+    try:
+        attrs = h5f[name].attrs
+    except KeyError as e:
+        msg = "not found"
+        raise LH5DecodeError(msg, h5f, name) from e
+    except AttributeError as e:
+        msg = "missing 'datatype' attribute"
+        raise LH5DecodeError(msg, h5f, name) from e
+
+    lgdotype = datatype.datatype(attrs["datatype"])
+
+    # scalars are dim-0 datasets
+    if lgdotype is types.Scalar:
+        return None
+
+    # structs don't have rows
+    if lgdotype is types.Struct:
+        return None
+
+    # tables should have elements with all the same length
+    if lgdotype is types.Table:
+        # read out each of the fields
+        rows_read = None
+        for field in datatype.get_struct_fields(attrs["datatype"]):
+            n_rows_read = read_n_rows(name + "/" + field, h5f)
+            if not rows_read:
+                rows_read = n_rows_read
+            elif rows_read != n_rows_read:
+                log.warning(
+                    f"'{field}' field in table '{name}' has {rows_read} rows, "
+                    f"{n_rows_read} was expected"
+                )
+        return rows_read
+
+    # length of vector of vectors is the length of its cumulative_length
+    if lgdotype is types.VectorOfVectors:
+        return read_n_rows(f"{name}/cumulative_length", h5f)
+
+    # length of vector of encoded vectors is the length of its decoded_size
+    if lgdotype in (types.VectorOfEncodedVectors, types.ArrayOfEncodedEqualSizedArrays):
+        return read_n_rows(f"{name}/encoded_data", h5f)
+
+    # return array length (without reading the array!)
+    if issubclass(lgdotype, types.Array):
+        # compute the number of rows to read
+        return h5f[name].shape[0]
+
+    msg = f"don't know how to read rows of LGDO {lgdotype.__name__}"
+    raise LH5DecodeError(msg, h5f, name)
+
+
+def get_h5_group(
+    group: str | h5py.Group,
+    base_group: h5py.Group,
+    grp_attrs: Mapping[str, Any] | None = None,
+    overwrite: bool = False,
+) -> h5py.Group:
+    """
+    Returns an existing :mod:`h5py` group from a base group or creates a
+    new one. Can also set (or replace) group attributes.
+
+    Parameters
+    ----------
+    group
+        name of the HDF5 group.
+    base_group
+        HDF5 group to be used as a base.
+    grp_attrs
+        HDF5 group attributes.
+    overwrite
+        whether overwrite group attributes, ignored if `grp_attrs` is
+        ``None``.
+    """
+    if not isinstance(group, h5py.Group):
+        if group in base_group:
+            group = base_group[group]
+        else:
+            group = base_group.create_group(group)
+            if grp_attrs is not None:
+                group.attrs.update(grp_attrs)
+            return group
+    if (
+        grp_attrs is not None
+        and len(set(grp_attrs.items()) ^ set(group.attrs.items())) > 0
+    ):
+        if not overwrite:
+            msg = (
+                f"Provided {grp_attrs=} are different from "
+                f"existing ones {dict(group.attrs)=} but overwrite flag is not set"
+            )
+            raise RuntimeError(msg)
+
+        log.debug(f"overwriting {group}.attrs...")
+        for key in group.attrs:
+            group.attrs.pop(key)
+        group.attrs.update(grp_attrs)
+
+    return group
 
 
 def expand_vars(expr: str, substitute: dict[str, str] | None = None) -> str:

lgdo/lh5_store.py

@@ -1,3 +1,8 @@
+"""
+.. warning::
+    This subpackage is deprecated, use :mod:`lgdo.lh5`.
+"""
+
 from __future__ import annotations
 
 import sys
@@ -23,12 +28,16 @@ from .types import (
     WaveformTable,  # noqa: F401
 )
 
-DEFAULT_HDF5_COMPRESSION = None
 LGDO = Union[Array, Scalar, Struct, VectorOfVectors]
-DEFAULT_HDF5_SETTINGS: dict[str, ...] = {"shuffle": True, "compression": "gzip"}
 
 
 class LH5Iterator(lh5.LH5Iterator):
+    """
+    .. warning::
+        This class is deprecated, use :class:`lgdo.lh5.iterator.LH5Iterator`.
+
+    """
+
     def __init__(
         self,
         lh5_files: str | list[str],
@@ -70,6 +79,11 @@ class LH5Iterator(lh5.LH5Iterator):
         write_start: int = 0,
         **h5py_kwargs,
     ) -> None:
+        """
+        .. warning::
+            This method is deprecated, use :meth:`lgdo.lh5.iterator.LH5Iterator.write`.
+
+        """
         warn(
             "lgdo.lh5_store has moved to a subfolder lgdo.lh5 containing LH5Iterator. "
             "The object you are calling this function from uses the old LH5Iterator class."
@@ -102,6 +116,11 @@ class LH5Iterator(lh5.LH5Iterator):
         obj_buf_start: int = 0,
         decompress: bool = True,
     ) -> tuple[LGDO, int]:
+        """
+        .. warning::
+            This method is deprecated, use :meth:`lgdo.lh5.iterator.LH5Iterator.read`.
+
+        """
         warn(
             "lgdo.lh5_store has moved to a subfolder lgdo.lh5 containing LH5Iterator. "
             "The object you are calling this function from uses the old LH5Iterator class."
@@ -124,6 +143,12 @@ class LH5Iterator(lh5.LH5Iterator):
 
 
 class LH5Store(lh5.LH5Store):
+    """
+    .. warning::
+        This class is deprecated, use :class:`lgdo.lh5.iterator.LH5Store`.
+
+    """
+
     def __init__(self, base_path: str = "", keep_open: bool = False):
         warn(
             "lgdo.lh5_store has moved to a subfolder lgdo.lh5 containing LH5Store. "
@@ -140,6 +165,11 @@ class LH5Store(lh5.LH5Store):
         lh5_file: str | h5py.File | list[str | h5py.File],
         **kwargs,
     ) -> tuple[LGDO, int]:
+        """
+        .. warning::
+            This method is deprecated, use :meth:`lgdo.lh5.store.LH5Store.read`.
+
+        """
         warn(
             "LH5Store.read_object() has been renamed to LH5Store.read(), "
             "Please update your code."
@@ -156,6 +186,11 @@ class LH5Store(lh5.LH5Store):
         lh5_file: str | h5py.File,
         **kwargs,
     ) -> tuple[LGDO, int]:
+        """
+        .. warning::
+            This method is deprecated, use :meth:`lgdo.lh5.store.LH5Store.write`.
+
+        """
         warn(
             "LH5Store.write_object() has been renamed to LH5Store.write(), "
             "Please update your code."
@@ -172,6 +207,12 @@ def load_dfs(
     lh5_group: str = "",
     idx_list: list[np.ndarray | list | tuple] | None = None,
 ) -> pd.DataFrame:
+    """
+    .. warning::
+        This function is deprecated, use :meth:`lgdo.types.lgdo.LGDO.view_as` to
+        view LGDO data as a Pandas data structure.
+
+    """
     warn(
         "lgdo.lh5_store has moved to a subfolder lgdo.lh5. "
         "Please replace 'from lgdo.lh5_store import load_dfs' with 'from lgdo.lh5 import load_dfs'. "
@@ -188,6 +229,12 @@ def load_nda(
     lh5_group: str = "",
     idx_list: list[np.ndarray | list | tuple] | None = None,
 ) -> dict[str, np.ndarray]:
+    """
+    .. warning::
+        This function is deprecated, use :meth:`lgdo.types.lgdo.LGDO.view_as` to
+        view LGDO data as a NumPy data structure.
+
+    """
     warn(
         "lgdo.lh5_store has moved to a subfolder lgdo.lh5. "
         "Please replace 'from lgdo.lh5_store import load_nda' with 'from lgdo.lh5 import load_nda'. "
@@ -199,6 +246,11 @@ def load_nda(
 
 
 def ls(lh5_file: str | h5py.Group, lh5_group: str = "") -> list[str]:
+    """
+    .. warning::
+        This function is deprecated, import :func:`lgdo.lh5.tools.ls`.
+
+    """
     warn(
         "lgdo.lh5_store has moved to a subfolder lgdo.lh5. "
         "Please replace 'from lgdo.lh5_store import ls' with 'from lgdo.lh5 import ls'. "
@@ -216,6 +268,11 @@ def show(
     indent: str = "",
     header: bool = True,
 ) -> None:
+    """
+    .. warning::
+        This function is deprecated, import :func:`lgdo.lh5.tools.show`.
+
+    """
     warn(
         "lgdo.lh5_store has moved to a subfolder lgdo.lh5. "
         "Please replace 'from lgdo.lh5_store import show' with 'from lgdo.lh5 import show'. "

lgdo/logging.py

@@ -1,4 +1,5 @@
 """This module implements some helpers for setting up logging."""
+
 from __future__ import annotations
 
 import logging
@@ -16,7 +17,7 @@ CRITICAL = logging.CRITICAL
 def setup(level: int = logging.INFO, logger: logging.Logger | None = None) -> None:
     """Setup a colorful logging output.
 
-    If `logger` is None, sets up only the ``pygama`` logger.
+    If `logger` is None, sets up only the ``lgdo`` logger.
 
     Parameters
     ----------
@@ -27,7 +28,7 @@ def setup(level: int = logging.INFO, logger: logging.Logger | None = None) -> No
 
     Examples
     --------
-    >>> from pygama import logging
+    >>> from lgdo import logging
     >>> logging.setup(level=logging.DEBUG)
     """
     handler = colorlog.StreamHandler()
@@ -36,7 +37,7 @@ def setup(level: int = logging.INFO, logger: logging.Logger | None = None) -> No
     )
 
     if logger is None:
-        logger = colorlog.getLogger("pygama")
+        logger = colorlog.getLogger("lgdo")
 
     logger.setLevel(level)
     logger.addHandler(handler)

lgdo/types/__init__.py

@@ -1,4 +1,5 @@
 """LEGEND Data Objects (LGDO) types."""
+
 from __future__ import annotations
 
 from .array import Array

lgdo/types/array.py

@@ -2,6 +2,7 @@
 Implements a LEGEND Data Object representing an n-dimensional array and
 corresponding utilities.
 """
+
 from __future__ import annotations
 
 import logging
@@ -182,6 +183,7 @@ class Array(LGDO):
             if self.nda.ndim == 1:
                 return pd.Series(self.nda, copy=False)
 
+            # if array is multi-dim, use awkward
             return akpd.from_awkward(self.view_as("ak"))
 
         if library == "np":
@@ -195,6 +197,7 @@ class Array(LGDO):
                 msg = "Pint does not support Awkward yet, you must view the data with_units=False"
                 raise ValueError(msg)
 
+            # NOTE: this is zero-copy!
             return ak.Array(self.nda)
 
         msg = f"{library} is not a supported third-party format."

lgdo/types/arrayofequalsizedarrays.py

@@ -2,6 +2,7 @@
 Implements a LEGEND Data Object representing an array of equal-sized arrays and
 corresponding utilities.
 """
+
 from __future__ import annotations
 
 from collections.abc import Iterator

lgdo/types/encoded.py

@@ -179,8 +179,8 @@ class VectorOfEncodedVectors(LGDO):
 
     def __str__(self) -> str:
         string = ""
-        pos = 0
-        for vec, size in self:
+        for pos, res in enumerate(self):
+            vec, size = res[0], res[1]
             if pos != 0:
                 string += " "
 
@@ -200,8 +200,6 @@ class VectorOfEncodedVectors(LGDO):
             if pos < len(self.encoded_data.cumulative_length):
                 string += ",\n"
 
-            pos += 1
-
         string = f"[{string}]"
 
         attrs = self.getattrs()
@@ -400,8 +398,7 @@ class ArrayOfEncodedEqualSizedArrays(LGDO):
 
     def __str__(self) -> str:
         string = ""
-        pos = 0
-        for vec in self:
+        for pos, vec in enumerate(self):
             if pos != 0:
                 string += " "
 
@@ -418,8 +415,6 @@ class ArrayOfEncodedEqualSizedArrays(LGDO):
             if pos < len(self.encoded_data.cumulative_length):
                 string += ",\n"
 
-            pos += 1
-
         string = f"[{string}] decoded_size={self.decoded_size}"
 
         attrs = self.getattrs()

lgdo/types/fixedsizearray.py

@@ -2,6 +2,7 @@
 Implements a LEGEND Data Object representing an n-dimensional array of fixed
 size and corresponding utilities.
 """
+
 from __future__ import annotations
 
 from typing import Any

lgdo/types/struct.py

@@ -2,6 +2,7 @@
 Implements a LEGEND Data Object representing a struct and corresponding
 utilities.
 """
+
 from __future__ import annotations
 
 import logging