fix: pull out .data from NumpyArray (#985)

* fix: pull out `.data` from `NumpyArray` * fix: catch `Content` objects
scikit-hep · Oct 12, 2023 · aff025f · aff025f
1 parent 576fbdb
commit aff025f
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git 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
@@ -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(
@@ -780,7 +782,7 @@ def extend(self, file, sink, data):
                             "how did this pass the type check?\n\n" + repr(content)
                         )
 
-                big_endian = numpy.asarray(content, dtype=datum["dtype"])
+                big_endian = numpy.asarray(content.data, dtype=datum["dtype"])
                 shape = tuple(shape) + big_endian.shape[1:]
 
                 if shape[1:] != datum["shape"]: