diff --git a/hepdata_lib/root_utils.py b/hepdata_lib/root_utils.py index ff1dd288..f4e435f7 100644 --- a/hepdata_lib/root_utils.py +++ b/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 diff --git a/tests/test_rootfilereader.py b/tests/test_rootfilereader.py index 23d1d947..02712886 100644 --- a/tests/test_rootfilereader.py +++ b/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