uproot-custom 1.1.3.dev1__tar.gz → 1.3.0__tar.gz

{uproot_custom-1.1.3.dev1 → uproot_custom-1.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: uproot-custom
-Version: 1.1.3.dev1
+Version: 1.3.0
 Summary: uproot extension for reading custom classes
 Author-Email: Mingrun Li <mrli@ihep.ac.cn>
 Classifier: Development Status :: 3 - Alpha

{uproot_custom-1.1.3.dev1 → uproot_custom-1.3.0}/cpp/include/uproot-custom/uproot-custom.hh

@@ -148,7 +148,9 @@ namespace uproot {
             if ( fBits & ( kIsReferenced ) ) skip( 2 ); // pidf
         }
 
+        const uint8_t* get_data() const { return m_data; }
         const uint8_t* get_cursor() const { return m_cursor; }
+        const uint32_t* get_offsets() const { return m_offsets; }
         const uint64_t entries() const { return m_entries; }
 
       private:

{uproot_custom-1.1.3.dev1 → uproot_custom-1.3.0}/cpp/src/uproot-custom.cc

@@ -378,39 +378,98 @@ namespace uproot {
 
     class CArrayReader : public IElementReader {
       public:
-        CArrayReader( std::string name, bool is_obj, const uint32_t flat_size,
+        CArrayReader( std::string name, bool is_obj, bool is_stdmap, const int64_t flat_size,
                       SharedReader element_reader )
             : IElementReader( name )
             , m_is_obj( is_obj )
+            , m_is_stdmap( is_stdmap )
             , m_flat_size( flat_size )
+            , m_offsets( std::make_shared<std::vector<uint32_t>>( 1, 0 ) )
             , m_element_reader( element_reader ) {}
 
         void read( BinaryBuffer& buffer ) override {
+
             PRINT_MSG( "CArrayReader::read() for " + m_name +
                        " with flat_size = " + std::to_string( m_flat_size ) +
                        ", is_obj = " + std::to_string( m_is_obj ) );
             PRINT_BUFFER( buffer );
-            if ( m_is_obj )
+
+            if ( m_flat_size > 0 )
             {
-                buffer.read_fNBytes();
-                buffer.read_fVersion();
+                if ( m_is_obj )
+                {
+                    buffer.read_fNBytes();
+                    buffer.read_fVersion();
+                    if ( m_is_stdmap ) buffer.skip( 6 );
+                }
+
+                for ( auto i = 0; i < m_flat_size; i++ )
+                {
+                    PRINT_MSG( "CArrayReader::read() reading element " + std::to_string( i ) );
+                    PRINT_BUFFER( buffer );
+                    m_element_reader->read( buffer );
+                }
+                PRINT_MSG( "" );
+                PRINT_MSG( "" );
             }
-            for ( auto i = 0; i < m_flat_size; i++ )
+
+            else
             {
-                PRINT_MSG( "CArrayReader::read() reading element " + std::to_string( i ) );
-                PRINT_BUFFER( buffer );
-                m_element_reader->read( buffer );
+                // get end-position
+                auto n_entries     = buffer.entries();
+                auto start_pos     = buffer.get_data();
+                auto entry_offsets = buffer.get_offsets();
+                auto cursor_pos    = buffer.get_cursor();
+                auto entry_end = std::find_if( entry_offsets, entry_offsets + n_entries + 1,
+                                               [start_pos, cursor_pos]( uint32_t offset ) {
+                                                   return start_pos + offset > cursor_pos;
+                                               } );
+
+                PRINT_MSG( "CArrayReader::read() cursor_pos = " +
+                           std::to_string( cursor_pos - start_pos ) +
+                           "entry_end = " + std::to_string( *entry_end ) );
+
+                if ( m_is_obj )
+                {
+                    buffer.read_fNBytes();
+                    buffer.read_fVersion();
+                    // if ( m_is_stdmap ) buffer.skip( 6 ); // Even std::map has no 6 bytes here.
+                }
+
+                uint32_t count = 0;
+                while ( buffer.get_cursor() < start_pos + *entry_end )
+                {
+
+                    PRINT_MSG( "CArrayReader::read() reading element " +
+                               std::to_string( count ) );
+                    PRINT_BUFFER( buffer );
+                    m_element_reader->read( buffer );
+                    count += 1;
+                }
+
+                PRINT_MSG( "" );
+                PRINT_MSG( "" );
+
+                m_offsets->push_back( m_offsets->back() + count );
             }
-            PRINT_MSG( "" );
-            PRINT_MSG( "" );
         }
 
-        py::object data() const override { return m_element_reader->data(); }
+        py::object data() const override {
+            if ( m_flat_size > 0 ) return m_element_reader->data();
+            else
+            {
+                auto offsets_array = make_array( m_offsets );
+                auto elements_data = m_element_reader->data();
+                return py::make_tuple( offsets_array, elements_data );
+            }
+        }
 
       private:
         bool m_is_obj;
-        const uint32_t m_flat_size;
+        bool m_is_stdmap;
+        const int64_t m_flat_size;
         SharedReader m_element_reader;
+        SharedVector<uint32_t> m_offsets;
     };
 
     /*
@@ -481,7 +540,7 @@ namespace uproot {
         register_reader<BaseObjectReader, std::vector<SharedReader>>( m, "BaseObjectReader" );
         register_reader<ObjectHeaderReader, std::vector<SharedReader>>( m,
                                                                         "ObjectHeaderReader" );
-        register_reader<CArrayReader, bool, uint32_t, SharedReader>( m, "CArrayReader" );
+        register_reader<CArrayReader, bool, bool, int64_t, SharedReader>( m, "CArrayReader" );
         register_reader<EmptyReader>( m, "EmptyReader" );
     }
 

{uproot_custom-1.1.3.dev1 → uproot_custom-1.3.0}/example/gen-demo-data/include/TComplicatedSTL.hh

@@ -17,13 +17,13 @@ using namespace std;
 class TComplicatedSTL : public TObject {
 
   public:
-    TComplicatedSTL() : TObject() {}
-
-    void fill() {
+    TComplicatedSTL() : TObject() {
         // Initialize 1 basic type element
         for ( int i = 0; i < 5; i++ )
         {
+
             vector<int> vec_int;
+            map<int, double> map_int_double;
             list<int> list_int;
             set<int> set_int;
             unordered_set<int> uset_int;
@@ -33,10 +33,21 @@ class TComplicatedSTL : public TObject {
                 list_int.push_back( 10 * i + j );
                 set_int.insert( 10 * i + j );
                 uset_int.insert( 10 * i + j );
+                map_int_double[10 * i + j] = 0.1 * ( 10 * i + j );
             }
 
-            // sequence like containers
+            // std::array
+            m_arr_int[i]     = 100 + i;
             m_arr_vec_int[i] = vec_int;
+            m_arr_str[i]     = "Hello, " + to_string( i );
+
+            // c-style array
+            m_carr_int[i]            = 10 + i;
+            m_carr_vec_int[i]        = vec_int;
+            m_carr_map_int_double[i] = map_int_double;
+            m_carr_str[i]            = "World, " + to_string( i );
+
+            // sequence like containers
             m_vec_list_int.push_back( list_int );
             m_list_set_int.push_back( set_int );
             m_vec_uset_int.push_back( uset_int );
@@ -47,23 +58,40 @@ class TComplicatedSTL : public TObject {
             m_map_set_int[i]   = set_int;
             m_umap_uset_int[i] = uset_int;
 
-            /* ------------------------------------ */
-            // mapping<sequence<object>> like containers
-            vector<TComplicatedSTL> vec_obj;
-            list<TComplicatedSTL*> list_objptr;
-            for ( int j = 0; j < 3; j++ )
+            // nested containers
+            vector<list<set<int>>> vec_list_set_int;
+            for ( int j = 0; j < 2; j++ )
             {
-                vec_obj.emplace_back();
-                list_objptr.push_back( new TComplicatedSTL() );
+                list<set<int>> list_set_int;
+                for ( int k = 0; k < 2; k++ )
+                {
+                    set<int> set_int;
+                    for ( int l = 0; l < 3; l++ )
+                    { set_int.insert( 100 * i + 10 * j + 2 * k + l ); }
+                    list_set_int.push_back( set_int );
+                }
+                vec_list_set_int.push_back( list_set_int );
             }
-            m_map_vec_obj[i]     = vec_obj;
-            m_map_list_objptr[i] = list_objptr;
+            m_map_vec_list_set_int[i] = vec_list_set_int;
         }
     }
 
   private:
-    // sequence like containers
+    int m_marker{ 114514 }; // just a marker
+
+    // c-style array
+    int m_carr_int[5]{};
+    vector<int> m_carr_vec_int[5]{};
+    map<int, double> m_carr_map_int_double[5]{};
+    string m_carr_str[5]{};
+
+    // std::array
+    array<int, 5> m_arr_int;
     array<vector<int>, 5> m_arr_vec_int;
+    array<string, 5> m_arr_str;
+    // ROOT-6.32.02 does not support std::array of map
+
+    // sequence like containers
     vector<list<int>> m_vec_list_int;
     list<set<int>> m_list_set_int;
     vector<unordered_set<int>> m_vec_uset_int;
@@ -74,9 +102,8 @@ class TComplicatedSTL : public TObject {
     map<int, set<int>> m_map_set_int;
     unordered_map<int, unordered_set<int>> m_umap_uset_int;
 
-    // mapping<sequence<object>> like containers
-    map<int, vector<TComplicatedSTL>> m_map_vec_obj;
-    map<int, list<TComplicatedSTL*>> m_map_list_objptr;
+    // nested containers
+    map<int, vector<list<set<int>>>> m_map_vec_list_set_int;
 
     ClassDef( TComplicatedSTL, 1 );
-};
+};

{uproot_custom-1.1.3.dev1 → uproot_custom-1.3.0}/example/gen-demo-data/src/main.cc

@@ -18,7 +18,7 @@ int main() {
     {
         ovrd_steamer    = TOverrideStreamer( i );
         complicated_stl = TComplicatedSTL();
-        complicated_stl.fill();
+
         t.Fill();
     }
 

uproot_custom-1.3.0/tests/test_AsCustom.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+
+import uproot
+import uproot_custom
+
+
+uproot_custom.AsCustom.target_branches |= {
+    "/my_tree:my_obj/m_carr_vec_int[3]",
+    "/my_tree:my_obj/m_int",
+    "/my_tree:my_obj/m_carr_tstring[3]",
+    "/my_tree:my_obj/m_carr2d_vec_int[2][3]",
+    "/my_tree:my_obj/m_carr2d_tstring[2][3]",
+}
+
+
+def test_AsCustom():
+    f = uproot.open(Path(__file__).parent / "test-data-1.root")
+    tree = f["my_tree"]
+    tree.arrays()

uproot_custom-1.3.0/tests/test_AsGroupedMap.py

@@ -0,0 +1,29 @@
+from pathlib import Path
+
+import uproot
+
+import uproot_custom
+
+uproot_custom.AsCustom.target_branches |= {
+    "/my_tree:complicated_stl/m_arr_str[5]",
+    "/my_tree:complicated_stl/m_arr_vec_int[5]",
+    "/my_tree:complicated_stl/m_carr_vec_int[5]",
+    "/my_tree:complicated_stl/m_carr_str[5]",
+    "/my_tree:complicated_stl/m_vec_uset_int",
+    "/my_tree:complicated_stl/m_arr_map_int_double[5]",
+    "/my_tree:complicated_stl/m_carr_map_int_double[5]",
+}
+
+uproot_custom.AsGroupedMap.target_branches |= {
+    "/my_tree:complicated_stl/m_map_vec_int/m_map_vec_int.second",
+    "/my_tree:complicated_stl/m_umap_list_int/m_umap_list_int.second",
+    "/my_tree:complicated_stl/m_map_set_int/m_map_set_int.second",
+    "/my_tree:complicated_stl/m_umap_uset_int/m_umap_uset_int.second",
+    "/my_tree:complicated_stl/m_map_vec_list_set_int/m_map_vec_list_set_int.second",
+}
+
+
+def test_AsGroupedMap():
+    f = uproot.open(Path(__file__).parent / "test-data-2.root")
+    tree = f["my_tree/complicated_stl"]
+    tree.arrays()

uproot_custom-1.3.0/uproot_custom/AsBinary.py

@@ -0,0 +1,95 @@
+from __future__ import annotations
+
+import awkward
+import numpy
+import uproot
+import uproot.extras
+import uproot.interpretation
+
+
+class AsBinary(uproot.interpretation.Interpretation):
+    """
+    Return binary data of the ``TBasket``. Pass an instance of this class
+    to :ref:`uproot.behaviors.TBranch.TBranch.array` like this:
+
+    .. code-block:: python
+        binary_data = branch.array(interpretation=AsBinary())
+
+    """
+
+    @property
+    def cache_key(self) -> str:
+        return id(self)
+
+    def basket_array(
+        self,
+        data,
+        byte_offsets,
+        basket,
+        branch,
+        context,
+        cursor_offset,
+        library,
+        interp_options,
+    ):
+        if byte_offsets is not None:
+            counts = byte_offsets[1:] - byte_offsets[:-1]
+        else:
+            counts = None
+
+        if library.name == "ak":
+            awkward = uproot.extras.awkward()
+            if counts is not None:
+                return awkward.unflatten(data, counts)
+            else:
+                fSize = branch.streamer.member("fSize")
+                return awkward.from_numpy(data.reshape(-1, fSize))
+
+        elif library.name == "np":
+            if counts is not None:
+                assert (
+                    numpy.unique(counts[1:] - counts[:-1]).size == 1
+                ), "The byte offsets must be uniform for NumPy arrays."
+
+                bytes_per_event = counts[0]
+                return data.reshape(-1, bytes_per_event)
+            else:
+                fSize = branch.streamer.member("fSize")
+                return data.reshape(-1, fSize).view(">u1")
+        else:
+            raise ValueError(
+                f"Unsupported library: {library.name}, can only use 'ak' or 'np'."
+            )
+
+    def final_array(
+        self,
+        basket_arrays,
+        entry_start,
+        entry_stop,
+        entry_offsets,
+        library,
+        branch,
+        options,
+    ):
+        basket_entry_starts = numpy.array(entry_offsets[:-1])
+        basket_entry_stops = numpy.array(entry_offsets[1:])
+
+        basket_start_idx = numpy.where(basket_entry_starts <= entry_start)[0].max()
+        basket_end_idx = numpy.where(basket_entry_stops >= entry_stop)[0].min()
+
+        arr_to_concat = [basket_arrays[i] for i in range(basket_start_idx, basket_end_idx + 1)]
+
+        relative_entry_start = entry_start - basket_entry_starts[basket_start_idx]
+        relative_entry_stop = entry_stop - basket_entry_starts[basket_start_idx]
+
+        if library.name == "ak":
+            awkward = uproot.extras.awkward()
+            return awkward.concatenate(arr_to_concat)[relative_entry_start:relative_entry_stop]
+
+        elif library.name == "np":
+            return numpy.concatenate(arr_to_concat)[relative_entry_start:relative_entry_stop]
+
+        else:
+            raise ValueError(
+                f"Unsupported library: {library.name}, can only use 'ak' or 'np'."
+            )

{uproot_custom-1.1.3.dev1 → uproot_custom-1.3.0}/uproot_custom/AsCustom.py

@@ -1,44 +1,13 @@
 from __future__ import annotations
 
-import re
-
+import awkward as ak
 import numpy as np
 import uproot
 import uproot.behaviors.TBranch
 import uproot.interpretation
 
 from uproot_custom.readers import read_branch
-
-
-def regularize_object_path(object_path: str) -> str:
-    return re.sub(r";[0-9]+", r"", object_path)
-
-
-_title_has_dims = re.compile(r"^([^\[\]]*)(\[[^\[\]]+\])+")
-_item_dim_pattern = re.compile(r"\[([1-9][0-9]*)\]")
-_item_any_pattern = re.compile(r"\[(.*)\]")
-
-
-def get_dims_from_branch(
-    branch: uproot.behaviors.TBranch.TBranch,
-) -> tuple[tuple[int, ...], bool]:
-    leaf = branch.member("fLeaves")[0]
-    title = leaf.member("fTitle")
-
-    dims, is_jagged = (), False
-
-    m = _title_has_dims.match(title)
-    if m is not None:
-        dims = tuple(int(x) for x in re.findall(_item_dim_pattern, title))
-        if dims == () and leaf.member("fLen") > 1:
-            dims = (leaf.member("fLen"),)
-
-        if any(
-            _item_dim_pattern.match(x) is None for x in re.findall(_item_any_pattern, title)
-        ):
-            is_jagged = True
-
-    return dims, is_jagged
+from uproot_custom.utils import get_dims_from_branch, regularize_object_path
 
 
 class AsCustom(uproot.interpretation.Interpretation):
@@ -65,6 +34,7 @@ class AsCustom(uproot.interpretation.Interpretation):
         self._branch = branch
         self._context = context
         self._simplify = simplify
+        self._typename = None
 
         # simplify streamer information
         self.all_streamer_info: dict[str, list[dict]] = {}
@@ -100,12 +70,14 @@ class AsCustom(uproot.interpretation.Interpretation):
         """
         The name of the type of the interpretation.
         """
-        dims, is_jagged = get_dims_from_branch(self._branch)
-        typename = self._branch.streamer.typename
-        if dims:
-            for i in dims:
-                typename += f"[{i}]"
-        return typename
+        if self._typename is None:
+            dims, is_jagged = get_dims_from_branch(self._branch)
+            typename = self._branch.streamer.typename
+            if dims:
+                for i in dims:
+                    typename += f"[{i}]"
+            self._typename = typename
+        return self._typename
 
     @property
     def cache_key(self) -> str:
@@ -133,9 +105,6 @@ class AsCustom(uproot.interpretation.Interpretation):
         """
         Concatenate the arrays from the baskets and return the final array.
         """
-
-        awkward = uproot.extras.awkward()
-
         basket_entry_starts = np.array(entry_offsets[:-1])
         basket_entry_stops = np.array(entry_offsets[1:])
 
@@ -143,7 +112,7 @@ class AsCustom(uproot.interpretation.Interpretation):
         basket_end_idx = np.where(basket_entry_stops >= entry_stop)[0].min()
 
         arr_to_concat = [basket_arrays[i] for i in range(basket_start_idx, basket_end_idx + 1)]
-        tot_array = awkward.concatenate(arr_to_concat)
+        tot_array = ak.concatenate(arr_to_concat)
 
         relative_entry_start = entry_start - basket_entry_starts[basket_start_idx]
         relative_entry_stop = entry_stop - basket_entry_starts[basket_start_idx]
@@ -165,10 +134,18 @@ class AsCustom(uproot.interpretation.Interpretation):
 
         full_branch_path = regularize_object_path(branch.object_path)
 
+        if branch.streamer is None:
+            cls_streamer_info = {
+                "fName": branch.name,
+                "fTypeName": self.typename,
+            }
+        else:
+            cls_streamer_info = branch.streamer.all_members
+
         return read_branch(
             data,
             byte_offsets,
-            branch.streamer.all_members,
+            cls_streamer_info,
             self.all_streamer_info,
             full_branch_path,
         )

uproot_custom-1.3.0/uproot_custom/AsGroupedMap.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import uproot
+import uproot.behaviors.TBranch
+import uproot.behaviors.TBranchElement
+import uproot.interpretation
+
+from uproot_custom.AsCustom import AsCustom
+from uproot_custom.utils import get_map_key_val_typenames
+
+
+class AsGroupedMap(AsCustom):
+    target_branches: set[str] = set()
+
+    def __init__(
+        self,
+        branch: uproot.behaviors.TBranch.TBranch,
+        context: dict,
+        simplify: bool,
+    ):
+        AsCustom.__init__(self, branch, context, simplify)
+
+        # 1:vector, 2:list, 3:deque, 4:map, 5:set, 6:multimap, 7:multiset, 12:unordered_map
+        stl_type = branch.parent.streamer.stl_type
+        assert stl_type in (
+            4,
+            6,
+            12,
+        ), f"Only map and multimap are supported for STL grouped branches, but got {stl_type}."
+
+        key_type_name, val_type_name = get_map_key_val_typenames(
+            branch.parent.streamer.typename
+        )
+
+        if branch == branch.parent.branches[0]:
+            self._typename = key_type_name + "[]"
+
+        elif branch == branch.parent.branches[1]:
+            self._typename = val_type_name + "[]"
+
+        else:
+            raise ValueError(
+                f"Branch {branch.name} not found in its parent branch {branch.parent.name}."
+            )

{uproot_custom-1.1.3.dev1 → uproot_custom-1.3.0}/uproot_custom/__init__.py

@@ -3,7 +3,8 @@ import uproot.behaviors.TBranch
 import uproot.interpretation.identify
 
 from uproot_custom.AsBinary import AsBinary
-from uproot_custom.AsCustom import AsCustom, regularize_object_path
+from uproot_custom.AsCustom import AsCustom
+from uproot_custom.AsGroupedMap import AsGroupedMap
 from uproot_custom.readers import (
     BaseObjectReader,
     BaseReader,
@@ -22,11 +23,11 @@ from uproot_custom.readers import (
     reconstruct_array,
     registered_readers,
 )
+from uproot_custom.utils import regularize_object_path
 
 ##########################################################################################
 #                                       Wrappers
 ##########################################################################################
-_is_TBranchElement_branches_wrapped = False
 _is_uproot_interpretation_of_wrapped = False
 
 _uproot_interpretation_of = uproot.interpretation.identify.interpretation_of
@@ -38,6 +39,9 @@ def custom_interpretation_of(
     if not hasattr(branch, "parent"):
         return _uproot_interpretation_of(branch, context, simplify)
 
+    if AsGroupedMap.match_branch(branch, context, simplify):
+        return AsGroupedMap(branch, context, simplify)
+
     if AsCustom.match_branch(branch, context, simplify):
         return AsCustom(branch, context, simplify)
 

{uproot_custom-1.1.3.dev1 → uproot_custom-1.3.0}/uproot_custom/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '1.1.3.dev1'
-__version_tuple__ = version_tuple = (1, 1, 3, 'dev1')
+__version__ = version = '1.3.0'
+__version_tuple__ = version_tuple = (1, 3, 0)
 
-__commit_id__ = commit_id = 'gab4d2fe79'
+__commit_id__ = commit_id = 'ga387b49b1'

{uproot_custom-1.1.3.dev1 → uproot_custom-1.3.0}/uproot_custom/readers.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import re
 from typing import Union
 
 import awkward as ak
@@ -9,17 +8,15 @@ import awkward.index
 import numpy as np
 
 import uproot_custom.cpp
+from uproot_custom.utils import (
+    get_map_key_val_typenames,
+    get_sequence_element_typename,
+    get_top_type_name,
+)
 
 registered_readers: set[type["BaseReader"]] = set()
 
 
-def get_top_type_name(type_name: str) -> str:
-    if type_name.endswith("*"):
-        type_name = type_name[:-1].strip()
-    type_name = type_name.replace("std::", "").strip()
-    return type_name.split("<")[0]
-
-
 def gen_tree_config(
     cls_streamer_info: dict,
     all_streamer_info: dict,
@@ -273,18 +270,6 @@ stl_typenames = {
 
 
 class STLSeqReader(BaseReader):
-    @staticmethod
-    def get_sequence_element_typename(type_name: str) -> str:
-        """
-        Get the element type name of a vector type.
-
-        e.g. vector<vector<int>> -> vector<int>
-        """
-        type_name = (
-            type_name.replace("std::", "").replace("< ", "<").replace(" >", ">").strip()
-        )
-        return re.match(r"^(vector|array|list|set|unordered_set)<(.*)>$", type_name).group(2)
-
     @classmethod
     def gen_tree_config(
         cls,
@@ -299,7 +284,7 @@ class STLSeqReader(BaseReader):
 
         fName = cls_streamer_info["fName"]
         fTypeName = cls_streamer_info["fTypeName"]
-        element_type = cls.get_sequence_element_typename(fTypeName)
+        element_type = get_sequence_element_typename(fTypeName)
         element_info = {
             "fName": fName,
             "fTypeName": element_type,
@@ -356,18 +341,6 @@ class STLMapReader(BaseReader):
     This class reads std::map from a binary parser.
     """
 
-    @staticmethod
-    def get_map_key_val_typenames(type_name: str) -> tuple[str, str]:
-        """
-        Get the key and value type names of a map type.
-
-        e.g. map<int, vector<int>> -> (int, vector<int>)
-        """
-        type_name = (
-            type_name.replace("std::", "").replace("< ", "<").replace(" >", ">").strip()
-        )
-        return re.match(r"^(map|unordered_map|multimap)<(.*),(.*)>$", type_name).groups()[1:3]
-
     @classmethod
     def gen_tree_config(
         cls,
@@ -381,7 +354,7 @@ class STLMapReader(BaseReader):
             return None
 
         fTypeName = cls_streamer_info["fTypeName"]
-        key_type_name, val_type_name = cls.get_map_key_val_typenames(fTypeName)
+        key_type_name, val_type_name = get_map_key_val_typenames(fTypeName)
 
         fName = cls_streamer_info["fName"]
         key_info = {
@@ -679,30 +652,37 @@ class CArrayReader(BaseReader):
         item_path,
         called_from_top,
     ):
-        if cls_streamer_info.get("fArrayDim", 0) == 0:
+        fTypeName = cls_streamer_info.get("fTypeName", "")
+        if not fTypeName.endswith("[]") and cls_streamer_info.get("fArrayDim", 0) == 0:
             return None
 
         fName = cls_streamer_info["fName"]
-        fTypeName = cls_streamer_info["fTypeName"]
-        fArrayDim = cls_streamer_info["fArrayDim"]
-        fMaxIndex = cls_streamer_info["fMaxIndex"]
+
+        if fTypeName.endswith("[]"):
+            fArrayDim = -1
+            fMaxIndex = -1
+            flat_size = -1
+        else:
+            fArrayDim = cls_streamer_info["fArrayDim"]
+            fMaxIndex = cls_streamer_info["fMaxIndex"]
+            flat_size = np.prod(fMaxIndex[:fArrayDim])
 
         element_streamer_info = cls_streamer_info.copy()
         element_streamer_info["fArrayDim"] = 0
+        while fTypeName.endswith("[]"):
+            fTypeName = fTypeName[:-2]
+        element_streamer_info["fTypeName"] = fTypeName
 
         element_tree_config = gen_tree_config(
             element_streamer_info,
             all_streamer_info,
         )
 
-        flat_size = np.prod(fMaxIndex[:fArrayDim])
-        assert flat_size > 0, f"flatten_size should be greater than 0, but got {flat_size}"
+        assert flat_size != 0, f"flatten_size should cannot be 0."
 
         # c-type number or TArray
-        if (
-            top_type_name in BasicTypeReader.typenames
-            or top_type_name in TArrayReader.typenames
-        ):
+        top_type_name = get_top_type_name(fTypeName)
+        if top_type_name in BasicTypeReader.typenames or fTypeName in TArrayReader.typenames:
             return {
                 "reader": cls,
                 "name": fName,
@@ -728,13 +708,27 @@ class CArrayReader(BaseReader):
         # STL
         elif top_type_name in stl_typenames:
             element_tree_config["with_header"] = False
+
             is_obj = not called_from_top
             if cls_streamer_info.get("fType", 0) == 500:
                 is_obj = True
+
+            # when is a ragged array, vector/map will have a reader
+            element_reader = element_tree_config.get("reader", None)
+            if (
+                flat_size < 0
+                and element_reader is not None
+                and element_reader != BasicTypeReader
+            ):
+                is_obj = True
+
+            is_stdmap = top_type_name in ["map", "unordered_map", "multimap"]
+
             return {
                 "reader": cls,
                 "name": fName,
                 "is_obj": is_obj,
+                "is_stdmap": is_stdmap,
                 "flat_size": flat_size,
                 "element_reader": element_tree_config,
                 "fMaxIndex": fMaxIndex,
@@ -755,6 +749,7 @@ class CArrayReader(BaseReader):
         return uproot_custom.cpp.CArrayReader(
             tree_config["name"],
             tree_config["is_obj"],
+            tree_config.get("is_stdmap", False),
             tree_config["flat_size"],
             element_reader,
         )
@@ -765,19 +760,33 @@ class CArrayReader(BaseReader):
             return None
 
         element_tree_config = tree_config["element_reader"]
-        fMaxIndex = tree_config["fMaxIndex"]
-        fArrayDim = tree_config["fArrayDim"]
-        shape = [fMaxIndex[i] for i in range(fArrayDim)]
-
-        element_data = reconstruct_array(
-            raw_data,
-            element_tree_config,
-        )
-
-        for s in shape[::-1]:
-            element_data = awkward.contents.RegularArray(element_data, int(s))
-
-        return element_data
+        flat_size = tree_config["flat_size"]
+
+        if flat_size > 0:
+            fMaxIndex = tree_config["fMaxIndex"]
+            fArrayDim = tree_config["fArrayDim"]
+            shape = [fMaxIndex[i] for i in range(fArrayDim)]
+
+            element_data = reconstruct_array(
+                raw_data,
+                element_tree_config,
+            )
+
+            for s in shape[::-1]:
+                element_data = awkward.contents.RegularArray(element_data, int(s))
+
+            return element_data
+
+        else:  # ragged array
+            offsets, element_raw_data = raw_data
+            element_data = reconstruct_array(
+                element_raw_data,
+                element_tree_config,
+            )
+            return ak.contents.ListOffsetArray(
+                ak.index.Index64(offsets),
+                element_data,
+            )
 
 
 class BaseObjectReader(BaseReader):

uproot_custom-1.3.0/uproot_custom/utils.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import re
+
+import uproot.behaviors.TBranch
+
+
+def regularize_object_path(object_path: str) -> str:
+    return re.sub(r";[0-9]+", r"", object_path)
+
+
+_title_has_dims = re.compile(r"^([^\[\]]*)(\[[^\[\]]+\])+")
+_item_dim_pattern = re.compile(r"\[([1-9][0-9]*)\]")
+_item_any_pattern = re.compile(r"\[(.*)\]")
+
+
+def get_dims_from_branch(
+    branch: uproot.behaviors.TBranch.TBranch,
+) -> tuple[tuple[int, ...], bool]:
+    leaf = branch.member("fLeaves")[0]
+    title = leaf.member("fTitle")
+
+    dims, is_jagged = (), False
+
+    m = _title_has_dims.match(title)
+    if m is not None:
+        dims = tuple(int(x) for x in re.findall(_item_dim_pattern, title))
+        if dims == () and leaf.member("fLen") > 1:
+            dims = (leaf.member("fLen"),)
+
+        if any(
+            _item_dim_pattern.match(x) is None for x in re.findall(_item_any_pattern, title)
+        ):
+            is_jagged = True
+
+    return dims, is_jagged
+
+
+def get_top_type_name(type_name: str) -> str:
+    if type_name.endswith("*"):
+        type_name = type_name[:-1].strip()
+    type_name = type_name.replace("std::", "").strip()
+    return type_name.split("<")[0]
+
+
+def get_sequence_element_typename(type_name: str) -> str:
+    """
+    Get the element type name of a vector type.
+
+    e.g. vector<vector<int>> -> vector<int>
+    """
+    type_name = type_name.replace("std::", "").replace("< ", "<").replace(" >", ">").strip()
+    return re.match(r"^(vector|array|list|set|unordered_set)<(.*)>$", type_name).group(2)
+
+
+def get_map_key_val_typenames(type_name: str) -> tuple[str, str]:
+    """
+    Get the key and value type names of a map type.
+
+    e.g. map<int, vector<int>> -> (int, vector<int>)
+    """
+    type_name = type_name.replace("std::", "").replace("< ", "<").replace(" >", ">").strip()
+    return re.match(r"^(map|unordered_map|multimap)<(.*),(.*)>$", type_name).groups()[1:3]

uproot_custom-1.1.3.dev1/tests/test_AsCustom.py

@@ -1,32 +0,0 @@
-from pathlib import Path
-
-import uproot
-import uproot_custom
-
-
-uproot_custom.AsCustom.target_branches |= {
-    "/my_tree:my_obj/m_carr_vec_int[3]",
-    "/my_tree:my_obj/m_int",
-    "/my_tree:my_obj/m_carr_tstring[3]",
-    "/my_tree:my_obj/m_carr2d_vec_int[2][3]",
-    "/my_tree:my_obj/m_carr2d_tstring[2][3]",
-    "/my_tree:complicated_stl/m_arr_vec_int[5]",
-    "/my_tree:complicated_stl/m_vec_uset_int",
-    "/my_tree:complicated_stl/m_vec_list_int",
-    "/my_tree:complicated_stl/m_list_set_int",
-}
-
-
-def test_AsCustom_1():
-    f = uproot.open(Path(__file__).parent / "test-data-1.root")
-    tree = f["my_tree"]
-    arr = tree.arrays()
-
-
-def test_AsCustom_2():
-    f = uproot.open(Path(__file__).parent / "test-data-2.root")
-    tree = f["my_tree"]
-    tree["complicated_stl/m_arr_vec_int[5]"].array()
-    tree["complicated_stl/m_vec_uset_int"].array()
-    tree["complicated_stl/m_vec_list_int"].array()
-    tree["complicated_stl/m_list_set_int"].array()

uproot_custom-1.1.3.dev1/uproot_custom/AsBinary.py

@@ -1,62 +0,0 @@
-from __future__ import annotations
-
-import awkward
-import numpy
-import uproot
-import uproot.extras
-import uproot.interpretation
-
-
-class AsBinary(uproot.interpretation.Interpretation):
-    """
-    Return binary data of the ``TBasket``. Pass an instance of this class
-    to :ref:`uproot.behaviors.TBranch.TBranch.array` like this:
-
-    .. code-block:: python
-        binary_data = branch.array(interpretation=AsBinary())
-
-    """
-
-    @property
-    def cache_key(self) -> str:
-        return id(self)
-
-    def basket_array(
-        self,
-        data,
-        byte_offsets,
-        basket,
-        branch,
-        context,
-        cursor_offset,
-        library,
-        interp_options,
-    ):
-        counts = byte_offsets[1:] - byte_offsets[:-1]
-        return awkward.unflatten(data, counts)
-
-    def final_array(
-        self,
-        basket_arrays,
-        entry_start,
-        entry_stop,
-        entry_offsets,
-        library,
-        branch,
-        options,
-    ):
-        basket_entry_starts = numpy.array(entry_offsets[:-1])
-        basket_entry_stops = numpy.array(entry_offsets[1:])
-
-        basket_start_idx = numpy.where(basket_entry_starts <= entry_start)[0].max()
-        basket_end_idx = numpy.where(basket_entry_stops >= entry_stop)[0].min()
-
-        arr_to_concat = [basket_arrays[i] for i in range(basket_start_idx, basket_end_idx + 1)]
-
-        awkward = uproot.extras.awkward()
-        tot_array = awkward.concatenate(arr_to_concat)
-
-        relative_entry_start = entry_start - basket_entry_starts[basket_start_idx]
-        relative_entry_stop = entry_stop - basket_entry_starts[basket_start_idx]
-
-        return tot_array[relative_entry_start:relative_entry_stop]