legend-pydataobj 1.5.1__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

@@ -2,6 +2,7 @@
 .. warning::
     This subpackage is deprecated, use :mod:`lgdo.lh5`.
 """
+
 from __future__ import annotations
 
 import sys
@@ -27,15 +28,14 @@ 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__(
@@ -82,6 +82,7 @@ class LH5Iterator(lh5.LH5Iterator):
         """
         .. 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. "
@@ -118,6 +119,7 @@ class LH5Iterator(lh5.LH5Iterator):
         """
         .. 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. "
@@ -144,6 +146,7 @@ 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):
@@ -165,6 +168,7 @@ class LH5Store(lh5.LH5Store):
         """
         .. warning::
             This method is deprecated, use :meth:`lgdo.lh5.store.LH5Store.read`.
+
         """
         warn(
             "LH5Store.read_object() has been renamed to LH5Store.read(), "
@@ -185,6 +189,7 @@ class LH5Store(lh5.LH5Store):
         """
         .. warning::
             This method is deprecated, use :meth:`lgdo.lh5.store.LH5Store.write`.
+
         """
         warn(
             "LH5Store.write_object() has been renamed to LH5Store.write(), "
@@ -206,6 +211,7 @@ def load_dfs(
     .. 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. "
@@ -227,6 +233,7 @@ def load_nda(
     .. 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. "
@@ -242,6 +249,7 @@ 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. "
@@ -263,6 +271,7 @@ def show(
     """
     .. warning::
         This function is deprecated, import :func:`lgdo.lh5.tools.show`.
+
     """
     warn(
         "lgdo.lh5_store has moved to a subfolder lgdo.lh5. "

lgdo/logging.py

@@ -1,4 +1,5 @@
 """This module implements some helpers for setting up logging."""
+
 from __future__ import annotations
 
 import logging

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

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

lgdo/types/table.py

@@ -2,6 +2,7 @@
 Implements a LEGEND Data Object representing a special struct of arrays of
 equal length and corresponding utilities.
 """
+
 from __future__ import annotations
 
 import logging
@@ -229,6 +230,29 @@ class Table(Struct):
         )
         return self.view_as(library="pd", cols=cols, prefix=prefix)
 
+    def flatten(self, _prefix="") -> Table:
+        """Flatten the table, if nested.
+
+        Returns a new :class:`Table` (that references, not copies, the existing
+        columns) with columns in nested tables being moved to the first level
+        (and renamed appropriately).
+
+        Examples
+        --------
+        >>> repr(tbl)
+        "Table(dict={'a': Array([1 2 3], attrs={'datatype': 'array<1>{real}'}), 'tbl': Table(dict={'b': Array([4 5 6], attrs={'datatype': 'array<1>{real}'}), 'tbl1': Table(dict={'z': Array([9 9 9], attrs={'datatype': 'array<1>{real}'})}, attrs={'datatype': 'table{z}'})}, attrs={'datatype': 'table{b,tbl1}'})}, attrs={'datatype': 'table{a,tbl}'})"
+        >>> tbl.flatten().keys()
+        dict_keys(['a', 'tbl__b', 'tbl__tbl1__z'])
+        """
+        flat_table = Table(size=self.size)
+        for key, obj in self.items():
+            if isinstance(obj, Table):
+                flat_table.join(obj.flatten(_prefix=f"{_prefix}{key}__"))
+            else:
+                flat_table.add_column(_prefix + key, obj)
+
+        return flat_table
+
     def eval(
         self,
         expr: str,
@@ -242,6 +266,13 @@ class Table(Struct):
         columns are viewed as :class:`ak.Array` and the respective routines are
         therefore available.
 
+        To columns nested in subtables can be accessed by scoping with two
+        underscores (``__``). For example: ::
+
+          tbl.eval("a + tbl2__b")
+
+        computes the sum of column `a` and column `b` in the subtable `tbl2`.
+
         Parameters
         ----------
         expr
@@ -285,15 +316,16 @@ class Table(Struct):
 
         # make a dictionary of low-level objects (numpy or awkward)
         # for later computation
+        flat_self = self.flatten()
         self_unwrap = {}
         has_ak = False
         for obj in c.co_names:
-            if obj in self.keys():
-                if isinstance(self[obj], VectorOfVectors):
-                    self_unwrap[obj] = self[obj].view_as("ak", with_units=False)
+            if obj in flat_self:
+                if isinstance(flat_self[obj], VectorOfVectors):
+                    self_unwrap[obj] = flat_self[obj].view_as("ak", with_units=False)
                     has_ak = True
                 else:
-                    self_unwrap[obj] = self[obj].view_as("np", with_units=False)
+                    self_unwrap[obj] = flat_self[obj].view_as("np", with_units=False)
 
         # use numexpr if we are only dealing with numpy data types
         if not has_ak:
@@ -319,7 +351,7 @@ class Table(Struct):
 
         # resort to good ol' eval()
         globs = {"ak": ak, "np": np}
-        out_data = eval(expr, globs, (self_unwrap | parameters))  # noqa: PGH001
+        out_data = eval(expr, globs, (self_unwrap | parameters))
 
         # need to convert back to LGDO
         if isinstance(out_data, ak.Array):