def preload_tree(self, uuid, treepath, tree):
"""To save a double-open when using NanoEventsFactory.from_file"""
key = "UprootSourceMapping:" + tuple_to_key((uuid, treepath))
self._cache.update(tree.file.array_cache)
self._cache.update(tree.file.object_cache)
tree.file.array_cache = self._cache
tree.file.object_cache = self._cache
self._cache[key] = tree
def _tree(self, uuid, treepath):
key = "UprootSourceMapping:" + tuple_to_key((uuid, treepath))
try:
return self._cache[key]
except KeyError:
pass
tree = self._fileopener.open_uuid(uuid)[treepath]
self._cache[key] = tree
return tree
def _column_source(self, uuid, path_in_source):
key = self.key_root() + tuple_to_key((uuid, path_in_source))
try:
return self._cache[key]
except KeyError:
pass
source = self._fileopener.open_uuid(uuid)[path_in_source]
self._cache[key] = source
return source
def events(self, runtime_cache=None):
mapping = CachedMapping(self._cc.storage,
UprootSourceMapping(self._cc))
partition_tuple = (
self._data["uuid"],
self._data["tree_name"],
"{0}-{1}".format(self._data["start"], self._data["stop"]),
)
factory = NanoEventsFactory(self._schema,
mapping,
tuple_to_key(partition_tuple),
cache=runtime_cache)
return factory.events()
def _tree(self, uuid, treepath):
key = "UprootSourceMapping:" + tuple_to_key((uuid, treepath))
try:
return self._cache[key]
except KeyError:
pass
pfn = self._uuid_pfnmap[uuid]
tree = uproot4.open(pfn + ":" + treepath, **self._uproot_options)
if str(tree.file.uuid) != uuid:
raise RuntimeError(
f"UUID of file {pfn} does not match expected value ({uuid})")
self._cache[key] = tree
return tree
def _from_mapping(
cls,
mapping,
partition_key,
base_form,
runtime_cache,
persistent_cache,
schemaclass,
metadata,
):
"""Quickly build NanoEvents from a root file
Parameters
----------
mapping : Mapping
The mapping of a column_source to columns.
partition_key : tuple
Basic information about the column source, uuid, paths.
base_form : dict
The awkward form describing the nanoevents interpretation of the mapped file.
runtime_cache : dict
A dict-like interface to a cache object. This cache is expected to last the
duration of the program only, and will be used to hold references to materialized
awkward arrays, etc.
persistent_cache : dict
A dict-like interface to a cache object. Only bare numpy arrays will be placed in this cache,
using globally-unique keys.
schemaclass : BaseSchema
A schema class deriving from `BaseSchema` and implementing the desired view of the file
metadata : dict
Arbitrary metadata to add to the `base.NanoEvents` object
"""
if persistent_cache is not None:
mapping = CachedMapping(persistent_cache, mapping)
if metadata is not None:
base_form["parameters"]["metadata"] = metadata
if not callable(schemaclass):
raise ValueError("Invalid schemaclass type")
schema = schemaclass(base_form)
if not isinstance(schema, BaseSchema):
raise RuntimeError("Invalid schema type")
return cls(schema,
mapping,
tuple_to_key(partition_key),
cache=runtime_cache)
def from_file(
cls,
file,
treepath="/Events",
entry_start=None,
entry_stop=None,
runtime_cache=None,
persistent_cache=None,
schemaclass=NanoAODSchema,
metadata=None,
uproot_options={},
access_log=None,
):
"""Quickly build NanoEvents from a file
Parameters
----------
file : str or uproot.reading.ReadOnlyDirectory
The filename or already opened file using e.g. ``uproot.open()``
treepath : str, optional
Name of the tree to read in the file
entry_start : int, optional
Start at this entry offset in the tree (default 0)
entry_stop : int, optional
Stop at this entry offset in the tree (default end of tree)
runtime_cache : dict, optional
A dict-like interface to a cache object. This cache is expected to last the
duration of the program only, and will be used to hold references to materialized
awkward1 arrays, etc.
persistent_cache : dict, optional
A dict-like interface to a cache object. Only bare numpy arrays will be placed in this cache,
using globally-unique keys.
schemaclass : BaseSchema
A schema class deriving from `BaseSchema` and implementing the desired view of the file
metadata : dict, optional
Arbitrary metadata to add to the `base.NanoEvents` object
uproot_options : dict, optional
Any options to pass to ``uproot.open``
access_log : list, optional
Pass a list instance to record which branches were lazily accessed by this instance
"""
if not issubclass(schemaclass, BaseSchema):
raise RuntimeError("Invalid schema type")
if isinstance(file, str):
tree = uproot.open(file, **uproot_options)[treepath]
elif isinstance(file, uproot.reading.ReadOnlyDirectory):
tree = file[treepath]
elif "<class 'uproot.rootio.ROOTDirectory'>" == str(type(file)):
raise RuntimeError(
"The file instance (%r) is an uproot3 type, but this module is only compatible with uproot4 or higher"
% file)
else:
raise TypeError("Invalid file type (%s)" % (str(type(file))))
if entry_start is None or entry_start < 0:
entry_start = 0
if entry_stop is None or entry_stop > tree.num_entries:
entry_stop = tree.num_entries
partition_tuple = (
str(tree.file.uuid),
tree.object_path,
"{0}-{1}".format(entry_start, entry_stop),
)
uuidpfn = {partition_tuple[0]: tree.file.file_path}
mapping = UprootSourceMapping(TrivialOpener(uuidpfn, uproot_options),
access_log=access_log)
mapping.preload_tree(partition_tuple[0], partition_tuple[1], tree)
if persistent_cache is not None:
mapping = CachedMapping(persistent_cache, mapping)
base_form = cls._extract_base_form(tree)
if metadata is not None:
base_form["parameters"]["metadata"] = metadata
schema = schemaclass(base_form)
return cls(schema,
mapping,
tuple_to_key(partition_tuple),
cache=runtime_cache)
def preload_tree(self, uuid, treepath, tree):
"""To save a double-open when using NanoEventsFactory.from_file"""
key = "UprootSourceMapping:" + tuple_to_key((uuid, treepath))
self._cache[key] = tree
def preload_column_source(self, uuid, path_in_source, source):
"""To save a double-open when using NanoEventsFactory.from_file"""
key = self.key_root() + tuple_to_key((uuid, path_in_source))
self._cache[key] = source
def from_file(
cls,
file,
treepath="/Events",
entry_start=None,
entry_stop=None,
runtime_cache=None,
persistent_cache=None,
schemaclass=NanoAODSchema,
metadata=None,
):
"""Quickly build NanoEvents from a file
Parameters
----------
file : str or uproot4.reading.ReadOnlyDirectory
The filename or already opened file using e.g. ``uproot4.open()``
treepath : str, optional
Name of the tree to read in the file
entry_start : int, optional
Start at this entry offset in the tree (default 0)
entry_stop : int, optional
Stop at this entry offset in the tree (default end of tree)
runtime_cache : dict, optional
A dict-like interface to a cache object. This cache is expected to last the
duration of the program only, and will be used to hold references to materialized
awkward1 arrays, etc.
persistent_cache : dict, optional
A dict-like interface to a cache object. Only bare numpy arrays will be placed in this cache,
using globally-unique keys.
schemaclass : BaseSchema
A schema class deriving from `BaseSchema` and implementing the desired view of the file
metadata : dict, optional
Arbitrary metadata to add to the `base.NanoEvents` object
"""
if not issubclass(schemaclass, BaseSchema):
raise RuntimeError("Invalid schema type")
if isinstance(file, str):
tree = uproot4.open(file + ":" + treepath)
elif isinstance(file, uproot4.reading.ReadOnlyDirectory):
tree = file[treepath]
if entry_start is None or entry_start < 0:
entry_start = 0
if entry_stop is None or entry_stop > tree.num_entries:
entry_stop = tree.num_entries
partition_tuple = (
str(tree.file.uuid),
tree.object_path,
"{0}-{1}".format(entry_start, entry_stop),
)
uuidpfn = {partition_tuple[0]: tree.file.file_path}
mapping = UprootSourceMapping(uuidpfn)
mapping.preload_tree(partition_tuple[0], partition_tuple[1], tree)
if persistent_cache is not None:
mapping = CachedMapping(persistent_cache, mapping)
base_form = cls._extract_base_form(tree)
if metadata is not None:
base_form["parameters"]["metadata"] = metadata
schema = schemaclass(base_form)
return cls(schema,
mapping,
tuple_to_key(partition_tuple),
cache=runtime_cache)