Retrieve histogram from stack via RootFileReader.retrieve_object

hepdata_lib/root_utils.py

@@ -101,7 +101,10 @@ def retrieve_object(self, path_to_object):
         try:
             obj = self.tfile.Get(parts[0])
             for part in parts[1:]:
-                obj = obj.GetPrimitive(part)
+                if isinstance(obj,r.THStack):  # pylint: disable=no-member
+                    obj = obj.GetHists().FindObject(part)
+                else:
+                    obj = obj.GetPrimitive(part)
 
             assert obj
 

tests/test_rootfilereader.py

@@ -17,6 +17,7 @@
 
 @pytest.mark.needs_root
 class TestRootFileReader(TestCase):
+    # pylint: disable=R0904
     """Test the RootFileReader class."""
 
     def test_tfile_setter(self):
@@ -784,6 +785,43 @@ def test_retrieve_object_canvas(self):
         # Clean up
         self.doCleanups()
 
+    def test_retrieve_object_stack(self):
+        '''Check that retrieve_object correctly reads from stack in canvas.'''
+        # Disable graphical output
+        ROOT.gROOT.SetBatch(ROOT.kTRUE)  # pylint: disable=no-member
+
+        # Create test histogram, plot on canvas, save to file
+        tfile = make_tmp_root_file(testcase=self)
+        histogram = ROOT.TH1D("testhist", "testhist", 10, 0, 1)  # pylint: disable=no-member
+        stack = ROOT.THStack("teststack","teststack")  # pylint: disable=no-member
+        stack.Add(histogram)
+        path_to_file = tfile.GetName()
+
+        canvas = ROOT.TCanvas()  # pylint: disable=no-member
+        stack.Draw("HIST")
+        canvas.Write("canvas")
+
+        reference = histogram.Clone("reference")
+        reference.SetDirectory(0)
+        if tfile:
+            tfile.Close()
+
+        # Read it back
+        reader = RootFileReader(path_to_file)
+        try:
+            readback = reader.retrieve_object("canvas/teststack/testhist")
+        except OSError:
+            print("RootFileReader.retrieve_object raised unexpected IOError!")
+            self.fail()
+
+        self.assertTrue(readback)
+        self.assertTrue(
+            histogram_compare_1d(reference, readback)
+            )
+
+        # Clean up
+        self.doCleanups()
+
     def test_retrieve_object_canvas_tpad(self):
         '''Check that retrieve_object correctly reads from canvas.'''
         # Disable graphical output