rootio: introduce a chain of Trees

rootio/tchain.go

@@ -0,0 +1,115 @@
+// Copyright 2018 The go-hep Authors.  All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package rootio
+
+type tchain struct {
+	trees []Tree
+}
+
+// Class returns the ROOT class of the argument.
+func (tchain) Class() string {
+	return "TChain"
+}
+
+// Name returns the name of the ROOT objet in the argument.
+func (t tchain) Name() string {
+	if len(t.trees) == 0 {
+		return ""
+	}
+	return t.trees[0].Name()
+}
+
+// Title returns the title of the ROOT object in the argument.
+func (t tchain) Title() string {
+	if len(t.trees) == 0 {
+		return ""
+	}
+	return t.trees[0].Title()
+}
+
+// Chain returns a tchain that is the concatenation of all the input Trees.
+func Chain(trees ...Tree) tchain {
+	if len(trees) == 0 {
+		return tchain{}
+	}
+	var t tchain
+	t.trees = make([]Tree, len(trees))
+	copy(t.trees, trees)
+	return t
+}
+
+// Entries returns the total number of entries.
+func (t tchain) Entries() int64 {
+	var v int64
+	for _, tree := range t.trees {
+		v += tree.Entries()
+	}
+	return v
+}
+
+// TotBytes return the total number of bytes before compression.
+func (t tchain) TotBytes() int64 {
+	var v int64
+	for _, tree := range t.trees {
+		v += tree.TotBytes()
+	}
+	return v
+}
+
+// ZipBytes returns the total number of bytes after compression.
+func (t tchain) ZipBytes() int64 {
+	var v int64
+	for _, tree := range t.trees {
+		v += tree.ZipBytes()
+	}
+	return v
+
+}
+
+// Branches returns the list of branches.
+func (t tchain) Branches() []Branch {
+	if len(t.trees) == 0 {
+		return nil
+	}
+	return t.trees[0].Branches()
+}
+
+// Branch returns the branch whose name is the argument.
+func (t tchain) Branch(name string) Branch {
+	if len(t.trees) == 0 {
+		return nil
+	}
+	return t.trees[0].Branch(name)
+}
+
+// Leaves returns direct pointers to individual branch leaves.
+func (t tchain) Leaves() []Leaf {
+	if len(t.trees) == 0 {
+		return nil
+	}
+	return t.trees[0].Leaves()
+}
+
+// getFile returns the underlying file.
+func (t tchain) getFile() *File {
+	if len(t.trees) == 0 {
+		return nil
+	}
+	return t.trees[0].getFile()
+}
+
+// loadEntry returns an error if there is a problem during the loading.
+func (t tchain) loadEntry(i int64) error {
+	if len(t.trees) == 0 {
+		return nil
+	}
+	return t.trees[0].loadEntry(i)
+}
+
+var (
+	_ Object = (*tchain)(nil)
+	_ Named  = (*tchain)(nil)
+	_ Tree   = (*tchain)(nil)
+)

rootio/tchain_test.go

@@ -0,0 +1,83 @@
+// Copyright 2018 The go-hep Authors.  All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package rootio_test
+
+import (
+	"testing"
+
+	"go-hep.org/x/hep/rootio"
+)
+
+func TestChain(t *testing.T) {
+	f, err := rootio.Open("testdata/chain.1.root")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+
+	obj, err := f.Get("tree")
+	if err != nil {
+		t.Fatal(err)
+	}
+	tree := obj.(rootio.Tree)
+	chain := rootio.Chain(tree)
+
+	if entry_tree, entry_chain := tree.Entries(), chain.Entries(); entry_tree != entry_chain {
+		t.Fatalf("entries differ: got=%v, want=%v\n", entry_chain, entry_tree)
+	}
+	if name_tree, name_chain := tree.Name(), chain.Name(); name_tree != name_chain {
+		t.Fatalf("names differ : got=%v, want=%v\n", name_chain, name_tree)
+	}
+	if title_tree, title_chain := tree.Title(), chain.Title(); title_tree != title_chain {
+		t.Fatalf("titles differ : got=%v, want=%v\n", title_chain, title_tree)
+	}
+}
+
+func TestChainEmpty(t *testing.T) {
+	chain := rootio.Chain()
+	if chain.Entries() != 0 {
+		t.Fatalf("unexpected entry : got=%v,want=0\n", chain.Entries())
+	}
+	if chain.Name() != "" {
+		t.Fatalf("unexpected name : got=%v, want empty string\n", chain.Name())
+	}
+	if chain.Title() != "" {
+		t.Fatalf("unexpected Title : got=%v, want empty string\n", chain.Title())
+	}
+}
+
+func TestTwoChains(t *testing.T) {
+	f1, err := rootio.Open("testdata/chain.1.root")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f1.Close()
+	f2, err := rootio.Open("testdata/chain.2.root")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f2.Close()
+
+	obj1, err1 := f1.Get("tree")
+	if err != nil {
+		t.Fatal(err1)
+	}
+	obj2, err2 := f2.Get("tree")
+	if err != nil {
+		t.Fatal(err2)
+	}
+
+	tree1 := obj1.(rootio.Tree)
+	tree2 := obj2.(rootio.Tree)
+	chain := rootio.Chain(tree1, tree2)
+
+	if entry_tree1, entry_tree2, entry_chain := tree1.Entries(), tree2.Entries(), chain.Entries(); entry_tree1+entry_tree2 != entry_chain {
+		t.Fatalf("entries differ: got=%v, want=%v\n", entry_chain, entry_tree1+entry_tree2)
+	}
+	if name_tree1, name_chain := tree1.Name(), chain.Name(); name_tree1 != name_chain {
+		t.Fatalf("names differ : got=%v, want=%v\n", name_chain, name_tree1)
+	}
+
+}