Updating the awkward typing and operation methods

scikit-hep · Nov 25, 2023 · e9868d8 · e9868d8
1 parent 1127f9e
commit e9868d8
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 11 deletions.
diff --git a/src/coffea/nanoevents/schemas/treemaker.py b/src/coffea/nanoevents/schemas/treemaker.py
@@ -173,7 +173,7 @@ def behavior(cls):
         return behavior
 
     @classmethod
-    def uproot_writeable(events):
+    def uproot_writeable(cls, events):
         """
         Converting a TreeMakerSchema event into something that is uproot
         writeable. Based off the discussion thread here [1], but added specific
@@ -185,15 +185,18 @@ def uproot_writeable(events):
         def _is_compat(a):
             """Is it a flat or 1-d jagged array?"""
             t = ak.type(a)
-            if isinstance(t, ak._ext.ArrayType):
-                if isinstance(t.type, ak._ext.PrimitiveType):
+            if isinstance(t, ak.types.ArrayType):
+                if isinstance(t._content, ak.types.NumpyType):
                     return True
-                if isinstance(t.type, ak._ext.ListType) and isinstance(
-                    t.type.type, ak._ext.PrimitiveType
+                if isinstance(t._content, ak.types.ListType) and isinstance(
+                    t._content._content, ak.types.NumpyType
                 ):
                     return True
             return False
 
+        def _make_packed(arr):
+            return ak.ak_to_packed.to_packed(ak.without_parameters(arr))
+
         def zip_composite(array):
             # Additional naming scheme to allow composite object read back
             _rename_lookup = {
@@ -207,7 +210,7 @@ def zip_composite(array):
             }
             return ak.zip(
                 {
-                    _rename_lookup.get(n, n): ak.packed(ak.without_parameters(array[n]))
+                    _rename_lookup.get(n, n): _make_packed(array[n])
                     for n in array.fields
                     if _is_compat(array[n])
                 }
@@ -229,10 +232,10 @@ def zip_composite(array):
                                 ak.flatten(events[bname][subname], axis=-1)
                             )
                         else:
-                            out[f"{bname}_{subname}"] = ak.flatten(
-                                events[bname][subname], axis=-1
+                            out[f"{bname}_{subname}"] = _make_packed(
+                                ak.flatten(events[bname][subname], axis=-1)
                             )
                 out[bname] = zip_composite(events[bname])
             else:
-                out[bname] = ak.packed(ak.without_parameters(events[bname]))
+                out[bname] = _make_packed(events[bname])
         return out
diff --git a/tests/test_nanoevents_treemaker.py b/tests/test_nanoevents_treemaker.py
@@ -105,6 +105,6 @@ def test_uproot_write():
     assert ak.all(orig_events.Jets.pt == test_events.Jets.pt)
     assert ak.all(orig_events.JetsAK8.x == test_events.JetsAK8.x)
     # Checking nested composite structure and their behavior
-    assert ak.all(orig_events.JetsAK8.subjet.pt == test_events.JetsAK8.subjet.pt)
-    assert ak.all(orig_events.JetsAK8.subjet.x == test_events.JetsAK8.subjet.x)
     assert ak.all(orig_events.Tracks.hitPattern == test_events.Tracks.hitPattern)
+    assert ak.all(orig_events.JetsAK8.subjets.pt == test_events.JetsAK8.subjets.pt)
+    assert ak.all(orig_events.JetsAK8.subjets.x == test_events.JetsAK8.subjets.x)