apply view to struc data if passed as unstruc #25

jacanchaplais · Jun 12, 2024 · b2c3569 · b2c3569
1 parent 1f31fcb
commit b2c3569
Showing 1 changed file with 12 additions and 19 deletions.
diff --git a/heparchy/write/hdf.py b/heparchy/write/hdf.py
@@ -5,7 +5,6 @@
 Provides the interface to write HEP data to the heparchy HDF5 format.
 """
 import functools as fn
-import itertools as it
 import typing as ty
 import warnings
 from enum import Enum
@@ -290,12 +289,10 @@ def edges(self) -> ty.NoReturn:
 
     @edges.setter
     def edges(self, data: base.VoidVector) -> None:
-        self._mk_dset(
-            name="edges",
-            data=data,
-            shape=data.shape,
-            dtype=np.dtype([("src", "<i4"), ("dst", "<i4")]),
-        )
+        dtype = np.dtype([("src", "<i4"), ("dst", "<i4")])
+        if not data.dtype.names:
+            data = data.view(dtype).reshape(-1)
+        self._mk_dset(name="edges", data=data, shape=data.shape, dtype=dtype)
         self._num_edges = len(data)
 
     @property
@@ -328,13 +325,11 @@ def pmu(self) -> ty.NoReturn:
 
     @pmu.setter
     def pmu(self, data: base.VoidVector) -> None:
+        dtype = np.dtype([(name, "<f8") for name in "xyze"])
+        if not data.dtype.names:
+            data = data.view(dtype).reshape(-1)
         self._set_num_pcls(data)
-        self._mk_dset(
-            name="pmu",
-            data=data,
-            shape=data.shape,
-            dtype=np.dtype(list(zip("xyze", it.repeat("<f8")))),
-        )
+        self._mk_dset(name="pmu", data=data, shape=data.shape, dtype=dtype)
 
     @property
     def color(self) -> ty.NoReturn:
@@ -343,13 +338,11 @@ def color(self) -> ty.NoReturn:
 
     @color.setter
     def color(self, data: base.AnyVector) -> None:
+        dtype = np.dtype([("color", "<i4"), ("anticolor", "<i4")])
+        if not data.dtype.names:
+            data = data.view(dtype).reshape(-1)
         self._set_num_pcls(data)
-        self._mk_dset(
-            name="color",
-            data=data,
-            shape=data.shape,
-            dtype=np.dtype([("color", "<i4"), ("anticolor", "<i4")]),
-        )
+        self._mk_dset(name="color", data=data, shape=data.shape, dtype=dtype)
 
     @property
     def pdg(self) -> ty.NoReturn: