From 85edc0f10ab4cfc73881be3e42e79ffd22e052c7 Mon Sep 17 00:00:00 2001 From: Angus Hollands Date: Thu, 12 Oct 2023 18:35:21 +0100 Subject: [PATCH] fix: catch `Content` objects --- src/uproot/_util.py | 14 ++++++++++---- src/uproot/writing/_cascadetree.py | 4 +++- 2 files changed, 13 insertions(+), 5 deletions(-) diff --git a/src/uproot/_util.py b/src/uproot/_util.py index 96d5bac92..b45b47ac0 100644 --- a/src/uproot/_util.py +++ b/src/uproot/_util.py @@ -77,12 +77,18 @@ def ensure_numpy(array, types=(numpy.bool_, numpy.integer, numpy.floating)): Returns an ``np.ndarray`` if ``array`` can be converted to an array of the desired type and raises TypeError if it cannot. """ + import uproot + + awkward = uproot.extras.awkward() with warnings.catch_warnings(): warnings.simplefilter("error", numpy.VisibleDeprecationWarning) - try: - out = numpy.asarray(array) - except (ValueError, numpy.VisibleDeprecationWarning) as err: - raise TypeError("cannot be converted to a NumPy array") from err + if isinstance(array, awkward.contents.Content): + out = awkward.to_numpy(array) + else: + try: + out = numpy.asarray(array) + except (ValueError, numpy.VisibleDeprecationWarning) as err: + raise TypeError("cannot be converted to a NumPy array") from err if not issubclass(out.dtype.type, types): raise TypeError(f"cannot be converted to a NumPy array of type {types}") return out diff --git a/src/uproot/writing/_cascadetree.py b/src/uproot/writing/_cascadetree.py index 231b68df7..80105f2c3 100644 --- a/src/uproot/writing/_cascadetree.py +++ b/src/uproot/writing/_cascadetree.py @@ -720,7 +720,9 @@ def extend(self, file, sink, data): ) ) else: - big_endian = numpy.asarray(branch_array, dtype=datum["dtype"]) + big_endian = uproot._util.ensure_numpy(branch_array).astype( + datum["dtype"] + ) if big_endian.shape != (len(branch_array),) + datum["shape"]: raise ValueError( "'extend' must fill branches with a consistent shape: has {}, trying to fill with {}".format(