Get all the field names and types from an RNTuple header.

Author: jpivarski

src/uproot/models/RNTuple.py

@@ -8,10 +8,20 @@
 
 import struct
 
+try:
+    import queue
+except ImportError:
+    import Queue as queue
+
 import uproot
 
 _rntuple_format1 = struct.Struct(">IIQIIQIIQ")
 
+# https://github.com/jblomer/root/blob/ntuple-binary-format-v1/tree/ntuple/v7/doc/specifications.md#envelopes
+_rntuple_frame_format = struct.Struct("<HHI")
+_rntuple_feature_flag_format = struct.Struct("<Q")
+_rntuple_num_bytes_fields = struct.Struct("<II")
+
 
 class Model_ROOT_3a3a_Experimental_3a3a_RNTuple(uproot.model.Model):
     """
@@ -40,6 +50,107 @@ def read_members(self, chunk, cursor, context, file):
             self._members["fReserved"],
         ) = cursor.fields(chunk, _rntuple_format1, context)
 
+        seek, nbytes = self._members["fSeekHeader"], self._members["fNBytesHeader"]
+        header_range = (seek, seek + nbytes)
+
+        seek, nbytes = self._members["fSeekFooter"], self._members["fNBytesFooter"]
+        footer_range = (seek, seek + nbytes)
+
+        notifications = queue.Queue()
+        compressed_header_chunk, compressed_footer_chunk = file.source.chunks(
+            [header_range, footer_range], notifications=notifications
+        )
+
+        if self._members["fNBytesHeader"] == self._members["fLenHeader"]:
+            self._header_chunk = compressed_header_chunk
+            self._header_cursor = uproot.source.cursor.Cursor(
+                self._members["fSeekHeader"]
+            )
+        else:
+            self._header_chunk = uproot.compression.decompress(
+                compressed_header_chunk,
+                uproot.source.cursor.Cursor(self._members["fSeekHeader"]),
+                context,
+                self._members["fNBytesHeader"],
+                self._members["fLenHeader"],
+            )
+            self._header_cursor = uproot.source.cursor.Cursor(0)
+
+        if self._members["fNBytesFooter"] == self._members["fLenFooter"]:
+            self._footer_chunk = compressed_footer_chunk
+            self._footer_cursor = uproot.source.cursor.Cursor(
+                self._members["fSeekFooter"]
+            )
+        else:
+            self._footer_chunk = uproot.compression.decompress(
+                compressed_footer_chunk,
+                uproot.source.cursor.Cursor(self._members["fSeekFooter"]),
+                context,
+                self._members["fNBytesFooter"],
+                self._members["fLenFooter"],
+            )
+            self._footer_cursor = uproot.source.cursor.Cursor(0)
+
+        self._header, self._footer = None, None
+
+    @property
+    def header(self):
+        if self._header is None:
+            cursor = self._header_cursor.copy()
+            context = {}
+
+            self._header = {}
+            self._header["frame"] = self._frame(self._header_chunk, cursor, context)
+
+            # https://github.com/jblomer/root/blob/ntuple-binary-format-v1/tree/ntuple/v7/doc/specifications.md#header-envelope
+            self._header["feature_flag"] = cursor.field(
+                self._header_chunk, _rntuple_feature_flag_format, context
+            )
+            self._header["name"] = cursor.rntuple_string(self._header_chunk, context)
+            self._header["description"] = cursor.rntuple_string(
+                self._header_chunk, context
+            )
+            self._header["author"] = cursor.rntuple_string(self._header_chunk, context)
+
+            cursor.skip(68)  # ???
+            num_fields_plus_one = cursor.field(
+                self._header_chunk, struct.Struct("<Q"), context
+            )
+
+            self._header["fields"] = [None] * (num_fields_plus_one - 1)
+            while any(x is None for x in self._header["fields"]):
+                field = {}
+
+                pos = cursor.index
+                field["num_bytes"] = cursor.field(
+                    self._header_chunk, struct.Struct("<I"), context
+                )
+                field["id"] = cursor.field(
+                    self._header_chunk, struct.Struct("<Q"), context
+                )
+                self._header["fields"][field["id"] - 1] = field
+
+                cursor.skip(48)  # ???
+                field["name"] = cursor.rntuple_string(self._header_chunk, context)
+                field["description"] = cursor.rntuple_string(
+                    self._header_chunk, context
+                )
+                field["type"] = cursor.rntuple_string(self._header_chunk, context)
+
+                cursor.move_to(pos + field["num_bytes"])
+
+        return self._header
+
+    @property
+    def footer(self):
+        raise NotImplementedError
+
+    def _frame(self, chunk, cursor, context):
+        version, min_version, num_bytes = cursor.fields(
+            chunk, _rntuple_frame_format, context
+        )
+        return {"version": version, "min_version": min_version, "num_bytes": num_bytes}
+
 
 uproot.classes[
     "ROOT::Experimental::RNTuple"

src/uproot/source/cursor.py

@@ -8,6 +8,7 @@
 
 from __future__ import absolute_import
 
+import datetime
 import struct
 import sys
 
@@ -22,6 +23,10 @@
 _raw_double32 = struct.Struct(">f")
 _raw_float16 = struct.Struct(">BH")
 
+# https://github.com/jblomer/root/blob/ntuple-binary-format-v1/tree/ntuple/v7/doc/specifications.md#basic-types
+_rntuple_string_length = struct.Struct("<I")
+_rntuple_datetime = struct.Struct("<Q")
+
 
 class Cursor(object):
     """
@@ -461,6 +466,20 @@ def classname(self, chunk, context, move=True):
         else:
             return out.decode(errors="surrogateescape")
 
+    def rntuple_string(self, chunk, context, move=True):
+        if move:
+            length = self.field(chunk, _rntuple_string_length, context)
+            return self.string_with_length(chunk, context, length)
+        else:
+            index = self._index
+            out = self.rntuple_string(chunk, context, move=True)
+            self._index = index
+            return out
+
+    def rntuple_datetime(self, chunk, context, move=True):
+        raw = self.field(chunk, _rntuple_datetime, context, move=move)
+        return datetime.datetime.fromtimestamp(raw)
+
     def debug(
         self,
         chunk,