The Context
My daily workflow largely consists of producing, styling, and circulating plots from a dataset to my advisor and collaborators. We use the C++ framework ROOT to generate and store histograms and I am writing my code in Python to take advantage of its Python bindings (PyROOT).
Since a ROOT file is the fundamental unit of our datasets, I wrote a simple context manager to facilitate the common task of opening a ROOT file, retrieving some histograms, and then closing the file.
import ROOT
class HistogramFile(object):
def __init__(self, filename):
self.filename = filename
def __enter__(self):
self.file = ROOT.TFile.Open(self.filename, 'read')
return self
def __exit__(self, exception_type, exception_value, traceback):
self.file.Close()
def get_histogram(self, name):
"""Return the histogram identified by name from the file.
"""
# The TFile::Get() method returns a pointer to an object stored in a ROOT file.
hist = self.file.Get(name)
if hist:
return hist
else:
raise RuntimeError('Unable to retrieve histogram named {0} from {1}'.format(name, self.filename))
This allows me to write the following snippet (imports of necessary modules implied)
f = ROOT.TFile.Open('dataset.root', 'read')
# Setup a canvas for plotting. The arguments are a name, an optional title, and the width and height in pixels.
canvas = ROOT.TCanvas('canvas', '', 500, 500)
hist = f.Get('electron_momentum')
hist.Draw()
canvas.SaveAs('plot.pdf')
f.Close()
in a more idiomatic fashion
with HistogramFile('dataset.root') as f:
canvas = ROOT.TCanvas('canvas', '', 500, 500)
hist = f.get_histogram('electron_momentum')
hist.Draw()
canvas.SaveAs('plot.pdf')
A dataset is often a collection of multiple ROOT files, so to make a plot I need to sum the histograms with the same name from each of its files together. The following snippet does not work, and demonstrated to me that the files must remain open in order to access their histograms.
with HistogramFile('dataset_part1.root') as f:
hist_1 = f.get_histogram('electron_momentum')
with HistogramFile('dataset_part2.root') as f:
hist_2 = f.get_histogram('electron_momentum')
# The next line causes a " *** Break *** segmentation violation"
hist_total = hist_1 + hist_2
After some searching, I discovered contextlib2's ExitStack, which allows me to programmatically handle a dynamic number of ROOT files.
import contextlib2
class Dataset(object):
def __init__(self, *filenames):
self.filenames = filenames
def __enter__(self):
with contextlib2.ExitStack() as stack:
self.files = [stack.enter_context(HistogramFile(fname) for fname in self.filenames]
self.close = stack.pop_all().close
return self
def __exit__(self, exception_type, exception_value, traceback):
self.close()
def get_histogram(self, name):
"""Return the sum of the histograms identified by name from all files.
"""
return sum(f.get_histogram(name) for f in self.files)
I now have a context manager of context managers that allows me to work with multiple files in a similar manner as before.
# A list of files that could be the result of globbing or os.listdir
dataset_files = ['dataset_part1.root', 'dataset_part2.root', 'dataset_part3.root']
with Dataset(*dataset_files) as dataset:
canvas = ROOT.TCanvas('canvas', '', 500, 500)
hist = dataset.get_histogram('electron_momentum')
hist.Draw()
canvas.SaveAs('plot.pdf')
The Issues
Are there problems with my code that I haven't addressed? Any subtleties or technicalities I have neglected?
I personally have not seen a context manager of context managers when perusing other people's code during my internet adventures. Is it suggestive of poor design choice? I mean, suppose I want to go one level higher in abstraction and design a new class which facilitates working with a collection of Dataset objects. How would I do that without making my code context managers all the way down?
ROOT
question here :) \$\endgroup\$with
open multiple files in parallel and extract the same tree from each of them: gist.github.com/graipher/9e8ad0776b6ab5a3935c4faf0f5fa091 \$\endgroup\$