Source code for rootpy.context

from __future__ import absolute_import

import os
from contextlib import contextmanager
# Note about locks: we don't need this in cases where ROOT has a
# thread-specific variable, so gDirectory and gPad are safe.
# Not so for gStyle, IsBatch and TH1.AddDirectory, so we use a lock in these
# cases. To prevent out-of-order lock grabbing, just use one reentrant lock for
# all of them.
import threading
LOCK = threading.RLock()

from . import ROOT
from . import log; log = log[__name__]

__all__ = [
    'preserve_current_style',
    'preserve_current_canvas',
    'preserve_current_directory',
    'preserve_batch_state',
    'invisible_canvas',
    'thread_specific_tmprootdir',
    'set_directory',
    'preserve_set_th1_add_directory',
    'working_directory',
    'do_nothing',
]


[docs]@contextmanager def preserve_current_style(): """ Context manager which ensures that the current style remains the current style when the context is left. """ # this should be 'Modern' by default with LOCK: old = ROOT.gStyle try: yield finally: old.cd()
[docs]@contextmanager def preserve_current_canvas(): """ Context manager which ensures that the current canvas remains the current canvas when the context is left. """ old = ROOT.gPad try: yield finally: if old: old.cd() elif ROOT.gPad: # Put things back how they were before. with invisible_canvas(): # This is a round-about way of resetting gPad to None. # No other technique I tried could do it. pass
[docs]@contextmanager def preserve_current_directory(): """ Context manager which ensures that the current directory remains the current directory when the context is left. """ old = ROOT.gDirectory try: yield finally: assert old, "BUG: assumptions were invalid. Please report this" # old is always valid and refers to ROOT.TROOT if no file is created. old.cd()
[docs]@contextmanager def preserve_batch_state(): """ Context manager which ensures the batch state is the same on exit as it was on entry. """ with LOCK: old = ROOT.gROOT.IsBatch() try: yield finally: ROOT.gROOT.SetBatch(old)
[docs]@contextmanager def invisible_canvas(): """ Context manager yielding a temporary canvas drawn in batch mode, invisible to the user. Original state is restored on exit. Example use; obtain X axis object without interfering with anything:: with invisible_canvas() as c: efficiency.Draw() g = efficiency.GetPaintedGraph() return g.GetXaxis() """ with preserve_current_canvas(): with preserve_batch_state(): ROOT.gROOT.SetBatch() c = ROOT.TCanvas() try: c.cd() yield c finally: c.Close() c.IsA().Destructor(c)
@contextmanager def thread_specific_tmprootdir(): """ Context manager which makes a thread specific gDirectory to avoid interfering with the current file. Use cases: A TTree Draw function which doesn't want to interfere with whatever gDirectory happens to be. Multi-threading where there are two threads creating objects with the same name which must reside in a directory. (again, this happens with TTree draw) """ with preserve_current_directory(): dname = "rootpy-tmp/thread/{0}".format( threading.current_thread().ident) d = ROOT.gROOT.mkdir(dname) if not d: d = ROOT.gROOT.GetDirectory(dname) assert d, "Unexpected failure, can't cd to tmpdir." d.cd() yield d @contextmanager def set_directory(robject): """ Context manager to temporarily set the directory of a ROOT object (if possible) """ if (not hasattr(robject, 'GetDirectory') or not hasattr(robject, 'SetDirectory')): log.warning("Cannot set the directory of a `{0}`".format( type(robject))) # Do nothing yield else: old_dir = robject.GetDirectory() try: robject.SetDirectory(ROOT.gDirectory) yield finally: robject.SetDirectory(old_dir) @contextmanager def preserve_set_th1_add_directory(state=True): """ Context manager to temporarily set TH1.AddDirectory() state """ with LOCK: status = ROOT.TH1.AddDirectoryStatus() try: ROOT.TH1.AddDirectory(state) yield finally: ROOT.TH1.AddDirectory(status) @contextmanager def working_directory(path): """ A context manager that changes the working directory to the given path, and then changes it back to its previous value on exit. """ prev_cwd = os.getcwd() os.chdir(path) try: yield finally: os.chdir(prev_cwd) @contextmanager def do_nothing(*args, **kwargs): """ A context manager that does... nothing! """ yield