Source code for

This module enhances IO-related ROOT functionality
from __future__ import absolute_import

import os
import re
import tempfile
from fnmatch import fnmatch
from collections import defaultdict

from .. import ROOT
from .. import asrootpy, QROOT
from ..base import Object, NamedObject
from ..decorators import snake_case_methods
from ..context import preserve_current_directory
from ..utils.path import expand as expand_path
from ..memory.keepalive import keepalive
from ..extern.shortuuid import uuid
from ..extern.six import string_types

__all__ = [

VALIDPATH = '^(?P<file>.+.root)(?:[/](?P<path>.+))?$'

[docs]class DoesNotExist(Exception): """ This exception is raised if an attempt is made to access an object that does not exist in a directory. """ pass
def autovivitree(): # return defaultdict(autovivitree) def splitfile(path): filename, _, path = path.partition(':' + os.path.sep) return filename, os.path.sep + path def wrap_path_handling(f): def get(self, name, *args, **kwargs): _name = os.path.normpath(name) if _name == '.': return self if _name == '..': return self._parent try: dirpath, _, path = _name.partition(os.path.sep) if path: if dirpath == '..': return self._parent.Get(path, *args, **kwargs) else: _dir = self.Get(dirpath) if not isinstance(_dir, _DirectoryBase): raise DoesNotExist _dir._parent = self _dir._path = os.path.join(self._path, dirpath) thing = get(_dir, path, *args, **kwargs) else: thing = f(self, _name, *args, **kwargs) if isinstance(thing, _DirectoryBase): thing._parent = self if isinstance(thing, _DirectoryBase): if isinstance(self, File): thing._path = os.path.normpath( (':' + os.path.sep).join([self._path, _name])) else: thing._path = os.path.normpath( os.path.join(self._path, _name)) return thing except DoesNotExist: raise DoesNotExist( "requested path '{0}' does not exist in {1}".format( name, self._path)) return get
[docs]def root_open(filename, mode=''): """ Open a ROOT file via ROOT's static ROOT.TFile.Open [1] function and return an asrootpy'd File. Parameters ---------- filename : string The absolute or relative path to the ROOT file. mode : string, optional (default='') Mode indicating how the file is to be opened. This can be either one of the options supported by ROOT.TFile.Open [2], or one of `a`, `a+`, `r`, `r+`, `w` or `w+`, with meanings as for the built-in `open()` function [3]. Returns ------- root_file : File an instance of rootpy's File subclass of ROOT's TFile. References ---------- .. [1] .. [2] .. [3] """ mode_map = {'a': 'UPDATE', 'a+': 'UPDATE', 'r': 'READ', 'r+': 'UPDATE', 'w': 'RECREATE', 'w+': 'RECREATE'} if mode in mode_map: mode = mode_map[mode] filename = expand_path(filename) prev_dir = ROOT.gDirectory root_file = ROOT.R.TFile.Open(filename, mode) if not root_file: raise IOError("could not open file: '{0}'".format(filename)) root_file.__class__ = File root_file._path = filename root_file._parent = root_file root_file._prev_dir = prev_dir # give Python ownership of the TFile so we can delete it ROOT.SetOwnership(root_file, True) return root_file
[docs]@snake_case_methods class Key(NamedObject, QROOT.TKey): """ A subclass of ROOT's TKey [1] References ---------- .. [1] """ _ROOT = QROOT.TKey
class _DirectoryBase(Object): def __str__(self): return "{0}('{1}')".format(self.__class__.__name__, self._path) def __repr__(self): return self.__str__() def __getattr__(self, attr): """ Natural naming support. Now you can get an object from a File/Directory with:: myfile.somedir.otherdir.histname """ # Be careful! If ``__getattr__`` ends up being called again here, # this can end up in an "infinite" recursion and stack overflow. # Directly call ROOT's Get() here since ``attr`` must anyway be a valid # identifier (not a path including subdirectories). thing = super(_DirectoryBase, self).Get(attr) if not thing: raise AttributeError( "{0} has no attribute '{1}'".format(self, attr)) thing = asrootpy(thing) if isinstance(thing, Directory): thing._path = os.path.join(self._path, thing.GetName()) thing._parent = self return thing def __setattr__(self, attr, value): """ Allow writing objects in a file with ``myfile.thing = myobject`` """ if (attr.startswith('_') # PyROOT also sets __lifeline or attr in self.__dict__ or not isinstance(value, ROOT.R.TObject)): # not fully initialized so use regular setattr return super(_DirectoryBase, self).__setattr__(attr, value) # pass to setitem self.__setitem__(attr, value) def __getitem__(self, name): return self.Get(name) def __setitem__(self, name, thing): """ Allow writing objects in a file with ``myfile['thing'] = myobject`` """ with preserve_current_directory(): thing.Write(name) def __iter__(self): return self.objects() def __enter__(self): curr_dir = ROOT.gDirectory if curr_dir != self: self._prev_dir = curr_dir return self def __exit__(self, type, value, traceback): self.Close() return False def cd_previous(self): """ cd to the gDirectory before this file was open. """ if self._prev_dir is None or isinstance(self._prev_dir, ROOT.TROOT): return False if isinstance(self._prev_dir, ROOT.TFile): if self._prev_dir.IsOpen() and self._prev_dir.IsWritable(): return True return False if not self._prev_dir.IsWritable(): # avoid warning from ROOT stating file is not writable return False prev_file = self._prev_dir.GetFile() if prev_file and prev_file.IsOpen(): return True return False def Close(self, *args): """ Like ROOT's Close but reverts to the gDirectory before this file was opened. """ super(_DirectoryBase, self).Close(*args) return self.cd_previous() def objects(self, cls=None): """ Return an iterater over all objects in this directory which are instances of `cls`. By default, iterate over all objects (`cls=None`). Parameters ---------- cls : a class, optional (default=None) If a class is specified, only iterate over objects that are instances of this class. Returns ------- A generator over the objects in this directory. Examples -------- $ rootpy browse myfile.root In [1]: list(f1.objects(R.Directory)) Out[1]: [Directory('mydirectory')] """ objs = (asrootpy(x.ReadObj(), warn=False) for x in self.GetListOfKeys()) if cls is not None: objs = (obj for obj in objs if isinstance(obj, cls)) return objs def keys(self, latest=False): """ Return a list of the keys in this directory. Parameters ---------- latest : bool, optional (default=False) If True then return a list of keys with unique names where only the key with the highest cycle number is included where multiple keys exist with the same name. Returns ------- keys : list List of keys """ if latest: keys = {} for key in self.keys(): name = key.GetName() if name in keys: if key.GetCycle() > keys[name].GetCycle(): keys[name] = key else: keys[name] = key return keys.values() return [asrootpy(key) for key in self.GetListOfKeys()] @wrap_path_handling def Get(self, path, rootpy=True, **kwargs): """ Return the requested object cast as its corresponding subclass in rootpy if one exists and ``rootpy=True``, otherwise return the unadulterated TObject. """ thing = super(_DirectoryBase, self).Get(path) if not thing: raise DoesNotExist # Ensure that the file we took the object from is alive at least as # long as the object being taken from it. # Note, Python does *not* own `thing`, it is ROOT's responsibility to # delete it in the C++ sense. (SetOwnership is False). However, ROOT # will delete the object when the TFile's destructor is run. # Therefore, when `thing` goes out of scope and the file referred to # by `this` has no references left, the file is destructed and calls # `thing`'s delete. # (this is thanks to the fact that weak referents (used by keepalive) # are notified when they are dead). keepalive(thing, self) if rootpy: return asrootpy(thing, **kwargs) return thing @wrap_path_handling def GetDirectory(self, path, rootpy=True, **kwargs): rdir = super(_DirectoryBase, self).GetDirectory(path) if not rdir: raise DoesNotExist if rootpy: return asrootpy(rdir, **kwargs) return rdir @wrap_path_handling def GetKey(self, path, cycle=9999, rootpy=True, **kwargs): """ Override TDirectory's GetKey and also handle accessing keys nested arbitrarily deep in subdirectories. """ key = super(_DirectoryBase, self).GetKey(path, cycle) if not key: raise DoesNotExist if rootpy: return asrootpy(key, **kwargs) return key def __contains__(self, path): """ Determine if a an object exists in the file at the path `path`:: if 'some/thing' in file: # do something """ try: self.GetKey(path) return True except DoesNotExist: return False def mkdir(self, path, title="", recurse=False): """ Make a new directory. If recurse is True, create parent directories as required. Return the newly created TDirectory. """ head, tail = os.path.split(os.path.normpath(path)) if tail == "": raise ValueError("invalid directory name: {0}".format(path)) with preserve_current_directory(): dest = self if recurse: parent_dirs = head.split(os.path.sep) for parent_dir in parent_dirs: try: newdest = dest.GetDirectory(parent_dir) dest = newdest except DoesNotExist: dest = dest.mkdir(parent_dir) elif head != "": dest = dest.GetDirectory(head) if tail in dest: raise ValueError("{0} already exists".format(path)) newdir = asrootpy(super(_DirectoryBase, dest).mkdir(tail, title)) return newdir def rm(self, path, cycle=';*'): """ Delete an object at `path` relative to this directory """ rdir = self with preserve_current_directory(): dirname, objname = os.path.split(os.path.normpath(path)) if dirname: rdir = rdir.Get(dirname) rdir.Delete(objname + cycle) # TODO: # def move(self, src, dest, newname=None): def copytree(self, dest_dir, src=None, newname=None, exclude=None, overwrite=False): """ Copy this directory or just one contained object into another directory. Parameters ---------- dest_dir : string or Directory The destination directory. src : string, optional (default=None) If ``src`` is None then this entire directory is copied recursively otherwise if ``src`` is a string path to an object relative to this directory, only that object will be copied. The copied object can optionally be given a ``newname``. newname : string, optional (default=None) An optional new name for the copied object. exclude : callable, optional (default=None) ``exclude`` can optionally be a function which takes ``(path, object_name)`` and if returns True excludes objects from being copied if the entire directory is being copied recursively. overwrite : bool, optional (default=False) If True, then overwrite existing objects with the same name. """ def copy_object(obj, dest, name=None): if name is None: name = obj.GetName() if not overwrite and name in dest: raise ValueError( "{0} already exists in {1} and `overwrite=False`".format( name, dest._path)) if isinstance(obj, ROOT.R.TTree): new_obj = obj.CloneTree(-1, "fast") new_obj.Write(name, ROOT.R.TObject.kOverwrite) else: obj.Write(name, ROOT.R.TObject.kOverwrite) with preserve_current_directory(): if isinstance(src, string_types): src = asrootpy(self.Get(src)) else: src = self if isinstance(dest_dir, string_types): try: dest_dir = asrootpy(self.GetDirectory(dest_dir)) except DoesNotExist: dest_dir = self.mkdir(dest_dir) if isinstance(src, ROOT.R.TDirectory): # Copy a directory cp_name = newname if newname is not None else src.GetName() # See if the directory already exists if cp_name not in dest_dir: # Destination directory doesn't exist, so make a new one new_dir = dest_dir.mkdir(cp_name) else: new_dir = dest_dir.get(cp_name) # Copy everything in the src directory to the destination for (path, dirnames, objects) in src.walk(maxdepth=0): # Copy all the objects for object_name in objects: if exclude and exclude(path, object_name): continue thing = src.Get(object_name) copy_object(thing, new_dir) for dirname in dirnames: if exclude and exclude(path, dirname): continue rdir = src.GetDirectory(dirname) # Recursively copy objects in subdirectories rdir.copytree( new_dir, exclude=exclude, overwrite=overwrite) else: # Copy an object copy_object(src, dest_dir, name=newname) def walk(self, top=None, path=None, depth=0, maxdepth=-1, class_ref=None, class_pattern=None, return_classname=False, treat_dirs_as_objs=False): """ Walk the directory structure and content in and below a directory. For each directory in the directory tree rooted at ``top`` (including ``top`` itself, but excluding '.' and '..'), yield a 3-tuple ``dirpath, dirnames, filenames``. Parameters ---------- top : string, optional (default=None) A path to a starting directory relative to this directory, otherwise start at this directory. path : string, optional (default=None) A path prepended as a prefix on the ``dirpath``. This argument is used internally as the recursion traverses down through subdirectories. depth : int, optional (default=0) The current depth, used internally as the recursion traverses down through subdirectories. max_depth : int, optional (default=-1) The maximum depth in the directory hierarchy to traverse. There is no limit applied by default. class_ref : class, optional (default=None) If not None then only include objects that are instances of ``class_ref``. class_pattern : string, optional (default=None) If not None then only include objects in ``filenames`` with class names that match ``class_pattern``. ``class_pattern`` should be a Unix shell-style wildcarded string. return_classname : bool, optional (default=False) If True, then each entry in ``filenames`` is a tuple of the form ``(filename, classname)``. treat_dirs_as_objs : bool, optional (default=False) If True, ``filenames`` contains directories as well. Returns ------- dirpath, dirnames, filenames : iterator An iterator over the 3-tuples ``dirpath, dirnames, filenames``. ``dirpath`` is a string, the path to the directory. ``dirnames`` is a list of the names of the subdirectories in ``dirpath`` (excluding '.' and '..'). ``filenames`` is a list of the names of the non-directory files/objects in ``dirpath``. Notes ----- The names in the lists are just names, with no path components. To get a full path (which begins with top) to a file or directory in ``dirpath``, use ``os.path.join(dirpath, name)``. """ dirnames, objectnames = [], [] tdirectory = self.GetDirectory(top) if top else self for key in tdirectory.keys(latest=True): name = key.GetName() classname = key.GetClassName() is_directory = classname.startswith('TDirectory') if is_directory: dirnames.append(name) if not is_directory or treat_dirs_as_objs: if class_ref is not None: tclass = ROOT.TClass.GetClass(classname, True, True) if not tclass or not tclass.InheritsFrom(class_ref.Class()): continue if class_pattern is not None: if not fnmatch(classname, class_pattern): continue name = (name if not return_classname else (name, classname)) objectnames.append(name) if path: dirpath = os.path.join(path, tdirectory.GetName()) elif not isinstance(tdirectory, ROOT.R.TFile): dirpath = tdirectory.GetName() else: dirpath = '' yield dirpath, dirnames, objectnames if depth == maxdepth: return for dirname in dirnames: rdir = tdirectory.GetDirectory(dirname) for x in rdir.walk( class_ref=class_ref, class_pattern=class_pattern, depth=depth + 1, maxdepth=maxdepth, path=dirpath, return_classname=return_classname, treat_dirs_as_objs=treat_dirs_as_objs): yield x
[docs]@snake_case_methods class Directory(_DirectoryBase, QROOT.TDirectoryFile): """ A subclass of ROOT's TDirectoryFile [1] References ---------- .. [1] """ _ROOT = QROOT.TDirectoryFile def __init__(self, name, title=None, classname='', parent=None): if title is None: title = name # grab previous directory before creating self self._prev_dir = ROOT.gDirectory self._parent = parent or self._prev_dir super(Directory, self).__init__(name, title, classname, parent or 0) self._post_init() def _post_init(self): self._path = self.GetName() # need to set _prev_dir here again if using rootpy.ROOT.TDirectory self._prev_dir = getattr(self, '_prev_dir', None) self._parent = getattr(self, '_parent', self._prev_dir)
class _FileBase(_DirectoryBase): def __init__(self, name, *args, **kwargs): # trigger finalSetup ROOT.R.kTRUE # grab previous directory before creating self self._prev_dir = ROOT.gDirectory super(_FileBase, self).__init__(name, *args, **kwargs) self._post_init() def _post_init(self): self._path = self.GetName() # need to set _prev_dir here again if using rootpy.ROOT.TFile self._prev_dir = getattr(self, '_prev_dir', None) self._parent = self def _populate_cache(self): """ Walk through the whole file and populate the cache all objects below the current path are added, i.e. for the contents with ina, inb and inab TH1F histograms:: /a/ina /b/inb /a/b/inab the cache is (omitting the directories):: cache[""]["obj"] = [("a", ("ina", "TH1F")), ("b", ("inb", "TH1F")), ("a/b", ("inab", "TH1F"))] ... cache[""]["a"]["b"]["obj"] = [("a/b", ("inab", "TH1F"))] """ self.cache = autovivitree() for path, dirs, objects in self.walk(return_classname=True, treat_dirs_as_objs=True): b = self.cache for d in ['']+path.split('/'): b = b[d] obj = [(path, o) for o in objects] if 'obj' in b: b['obj'] += obj else: b['obj'] = obj def find(self, regexp, negate_regexp=False, class_pattern=None,, refresh_cache=False): """ yield the full path of the matching regular expression and the match itself """ if refresh_cache or not hasattr(self, 'cache'): self._populate_cache() b = self.cache split_regexp = regexp.split('/') # traverse as deep as possible in the cache # special case if the first character is not the root, i.e. not "" if split_regexp[0] == '': for d in split_regexp: if d in b: b = b[d] else: break else: b = b[''] # perform the search for path, (obj, classname) in b['obj']: if class_pattern: if not fnmatch(classname, class_pattern): continue joined_path = os.path.join(*['/', path, obj]) result = find_fnc(regexp, joined_path) if (result is not None) ^ negate_regexp: yield joined_path, result
[docs]@snake_case_methods class File(_FileBase, QROOT.TFile): """ A subclass of ROOT's TFile [1] Examples -------- >>> from import File >>> from rootpy.testdata import get_filepath >>> f = File(get_filepath(), 'read') >>> list(f) [Directory('means'), Directory('scales'), Directory('gaps'), Directory('efficiencies'), Directory('dimensions'), Directory('graphs')] >>> f.means Directory('rootpy/testdata/test_file.root/means') References ---------- .. [1] """ _ROOT = QROOT.TFile # Override .Open open = staticmethod(root_open) Open = staticmethod(root_open)
[docs]@snake_case_methods class MemFile(_FileBase, QROOT.TMemFile): """ A subclass of ROOT's TMemFile [1] Examples -------- >>> from import MemFile >>> f = MemFile() References ---------- .. [1] """ _ROOT = QROOT.TMemFile def __init__(self, name=None, mode='recreate'): if name is None: name = '{0}_{1}'.format(self.__class__.__name__, uuid()) super(MemFile, self).__init__(name, mode)
[docs]@snake_case_methods class TemporaryFile(File): """ A temporary ROOT file that is automatically deleted when closed. Python's :func:`tempfile.mkstemp` [1] is used to obtain a temporary file in the most secure manner possible. Keyword arguments are passed directly to :func:`tempfile.mkstemp` [1] References ---------- .. [1] """ def __init__(self, suffix='.root', **kwargs): self.__fd, self.__tmp_path = tempfile.mkstemp(suffix=suffix, **kwargs) super(TemporaryFile, self).__init__(self.__tmp_path, 'recreate')
[docs] def Close(self): """ The physical file is automatically deleted after being closed. """ super(TemporaryFile, self).Close() os.close(self.__fd) os.remove(self.__tmp_path)