Skip to content

Commit

Permalink
Switched to using enums
Browse files Browse the repository at this point in the history
  • Loading branch information
@ariostas
ariostas committed Oct 4, 2024
1 parent 0104778 commit 1c028af
Show file tree
Hide file tree
Showing 2 changed files with 53 additions and 35 deletions.
60 changes: 39 additions & 21 deletions src/uproot/const.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
"""
from __future__ import annotations

from enum import IntEnum

import numpy

# determines when a file is "big"
Expand Down Expand Up @@ -208,31 +210,47 @@
"splitzigzagint16": 28,
}

rntuple_locator_type_uri = 0x01
rntuple_locator_type_daos = 0x02

rntuple_env_type_header = 0x01
rntuple_env_type_footer = 0x02
rntuple_env_type_pagelist = 0x03
rntuple_env_type_metadata = 0x04
class RNTupleLocatorType(IntEnum):
URI = 0x01
DAOS = 0x02


class RNTupleEnvelopeType(IntEnum):
HEADER = 0x01
FOOTER = 0x02
PAGELIST = 0x03
METADATA = 0x04


class RNTupleFieldRole(IntEnum):
LEAF = 0x00
VECTOR = 0x01
STRUCT = 0x02
UNION = 0x03
UNSPLIT = 0x04


class RNTupleFieldFlag(IntEnum):
REPETITIVE = 0x01
PROJECTED = 0x02
CHECKSUM = 0x04


class RNTupleColumnFlag(IntEnum):
DEFERRED = 0x08

rntuple_field_role_leaf = 0x00
rntuple_field_role_vector = 0x01
rntuple_field_role_struct = 0x02
rntuple_field_role_union = 0x03
rntuple_field_role_unsplit = 0x04

rntuple_field_flag_repetitive = 0x01
rntuple_field_flag_projected = 0x02
rntuple_field_flag_checksum = 0x04
class RNTupleExtraTypeIdentifier(IntEnum):
ROOT = 0x00

rntuple_col_flag_deferred = 0x08

rntuple_extra_type_identifier_root = 0x00
class RNTupleUserMetadataType(IntEnum):
INT = 0x01
BOOL = 0x02
DOUBLE = 0x03
STRING = 0x04

rntuple_user_metadata_type_int = 0x01
rntuple_user_metadata_type_bool = 0x02
rntuple_user_metadata_type_double = 0x03
rntuple_user_metadata_type_string = 0x04

rntuple_cluster_flag_sharded = 0x01
class RNTupleClusterFlag(IntEnum):
SHARDED = 0x01
28 changes: 14 additions & 14 deletions src/uproot/models/RNTuple.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -393,7 +393,7 @@ def field_form(self, this_id, seen):
seen.add(this_id)
structural_role = this_record.struct_role
if (
structural_role == uproot.const.rntuple_field_role_leaf
structural_role == uproot.const.RNTupleFieldRole.LEAF
and this_record.repetition == 0
):
# deal with std::atomic
Expand All @@ -408,7 +408,7 @@ def field_form(self, this_id, seen):
# base case of recursion
# n.b. the split may happen in column
return self.col_form(this_id)
elif structural_role == uproot.const.rntuple_field_role_leaf:
elif structural_role == uproot.const.RNTupleFieldRole.LEAF:
if this_id in self._related_ids:
# std::array has only one subfield
child_id = self._related_ids[this_id][0]
Expand All @@ -418,7 +418,7 @@ def field_form(self, this_id, seen):
inner = self.col_form(this_id)
keyname = f"RegularForm-{this_id}"
return ak.forms.RegularForm(inner, this_record.repetition, form_key=keyname)
elif structural_role == uproot.const.rntuple_field_role_vector:
elif structural_role == uproot.const.RNTupleFieldRole.VECTOR:
if this_id not in self._related_ids or len(self._related_ids[this_id]) != 1:
keyname = f"vector-{this_id}"
newids = self._related_ids.get(this_id, [])
Expand All @@ -442,15 +442,15 @@ def field_form(self, this_id, seen):
child_id = self._related_ids[this_id][0]
inner = self.field_form(child_id, seen)
return ak.forms.ListOffsetForm("i64", inner, form_key=keyname)
elif structural_role == uproot.const.rntuple_field_role_struct:
elif structural_role == uproot.const.RNTupleFieldRole.STRUCT:
newids = []
if this_id in self._related_ids:
newids = self._related_ids[this_id]
# go find N in the rest, N is the # of fields in struct
recordlist = [self.field_form(i, seen) for i in newids]
namelist = [field_records[i].field_name for i in newids]
return ak.forms.RecordForm(recordlist, namelist, form_key="whatever")
elif structural_role == uproot.const.rntuple_field_role_union:
elif structural_role == uproot.const.RNTupleFieldRole.UNION:
keyname = self.col_form(this_id)
newids = []
if this_id in self._related_ids:
Expand All @@ -460,7 +460,7 @@ def field_form(self, this_id, seen):
"i8", "i64", recordlist, form_key=keyname + "-union"
)
return ak.forms.IndexedOptionForm("i64", inner, form_key=keyname)
elif structural_role == uproot.const.rntuple_field_role_unsplit:
elif structural_role == uproot.const.RNTupleFieldRole.UNSPLIT:
raise NotImplementedError(
f"Unsplit fields are not supported. {this_record}"
)
Expand Down Expand Up @@ -724,7 +724,7 @@ def read(self, chunk, cursor, context):
out = MetaData(type(self).__name__)
out.env_header = _envelop_header(chunk, cursor, context)
assert (
out.env_header["env_type_id"] == uproot.const.rntuple_env_type_pagelist
out.env_header["env_type_id"] == uproot.const.RNTupleEnvelopeType.PAGELIST
), f"env_type_id={out.env_header['env_type_id']}"
out.header_checksum = cursor.field(chunk, _rntuple_checksum_format, context)
out.cluster_summaries = self.list_cluster_summaries.read(chunk, cursor, context)
Expand Down Expand Up @@ -813,17 +813,17 @@ def read(self, chunk, cursor, context):
out.struct_role,
out.flags,
) = cursor.fields(chunk, _rntuple_field_description_format, context)
if out.flags == uproot.const.rntuple_field_flag_repetitive:
if out.flags == uproot.const.RNTupleFieldFlag.REPETITIVE:
out.repetition = cursor.field(chunk, _rntuple_repetition_format, context)
out.source_field_id = None
out.checksum = None
elif out.flags == uproot.const.rntuple_field_flag_projected:
elif out.flags == uproot.const.RNTupleFieldFlag.PROJECTED:
out.repetition = 0
out.source_field_id = cursor.field(
chunk, _rntuple_source_field_id_format, context
)
out.checksum = None
elif out.flags == uproot.const.rntuple_field_flag_checksum:
elif out.flags == uproot.const.RNTupleFieldFlag.CHECKSUM:
out.repetition = 0
out.source_field_id = None
out.checksum = cursor.field(chunk, _rntuple_checksum_format, context)
Expand All @@ -844,7 +844,7 @@ def read(self, chunk, cursor, context):
out.type, out.nbits, out.field_id, out.flags, out.repr_idx = cursor.fields(
chunk, _rntuple_column_record_format, context
)
if out.flags & uproot.const.rntuple_col_flag_deferred:
if out.flags & uproot.const.RNTupleColumnFlag.DEFERRED:
out.first_element_index = cursor.field(
chunk, _rntuple_first_element_index_format, context
)
Expand Down Expand Up @@ -893,7 +893,7 @@ def read(self, chunk, cursor, context):
out = MetaData(type(self).__name__)
out.env_header = _envelop_header(chunk, cursor, context)
assert (
out.env_header["env_type_id"] == uproot.const.rntuple_env_type_header
out.env_header["env_type_id"] == uproot.const.RNTupleEnvelopeType.HEADER
), f"env_type_id={out.env_header['env_type_id']}"
out.feature_flag = cursor.field(chunk, _rntuple_feature_flag_format, context)
out.name, out.ntuple_description, out.writer_identifier = (
Expand Down Expand Up @@ -937,7 +937,7 @@ def read(self, chunk, cursor, context):
)
out.flags = out.num_entries >> 28
out.num_entries &= 0x0FFFFFFF
if out.flags == uproot.const.rntuple_cluster_flag_sharded:
if out.flags == uproot.const.RNTupleClusterFlag.SHARDED:
out.col_group_id = cursor.field(
chunk, _rntuple_col_group_id_format, context
)
Expand Down Expand Up @@ -994,7 +994,7 @@ def read(self, chunk, cursor, context):
out = MetaData("Footer")
out.env_header = _envelop_header(chunk, cursor, context)
assert (
out.env_header["env_type_id"] == uproot.const.rntuple_env_type_footer
out.env_header["env_type_id"] == uproot.const.RNTupleEnvelopeType.FOOTER
), f"env_type_id={out.env_header['env_type_id']}"
out.feature_flag = cursor.field(chunk, _rntuple_feature_flag_format, context)
out.header_checksum = cursor.field(chunk, _rntuple_checksum_format, context)
Expand Down

0 comments on commit 1c028af

Please sign in to comment.