go-hep.org/x/hep@v0.38.1/groot/rcmd/copy.go (about)

     1  // Copyright ©2020 The go-hep Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package rcmd
     6  
     7  import (
     8  	"fmt"
     9  	"log"
    10  	stdpath "path"
    11  	"regexp"
    12  
    13  	"go-hep.org/x/hep/groot"
    14  	"go-hep.org/x/hep/groot/riofs"
    15  	"go-hep.org/x/hep/groot/root"
    16  	"go-hep.org/x/hep/groot/rtree"
    17  )
    18  
    19  // Copy copies the content of the ROOT files fnames into the output
    20  // ROOT file named oname.
    21  func Copy(oname string, fnames []string) error {
    22  	o, err := groot.Create(oname)
    23  	if err != nil {
    24  		return fmt.Errorf("could not create output ROOT file %q: %w", oname, err)
    25  	}
    26  	defer o.Close()
    27  
    28  	var cmd copyCmd
    29  	for _, arg := range fnames {
    30  		err := cmd.process(o, arg)
    31  		if err != nil {
    32  			return err
    33  		}
    34  	}
    35  
    36  	err = o.Close()
    37  	if err != nil {
    38  		return fmt.Errorf("could not close output ROOT file %q: %w", oname, err)
    39  	}
    40  	return nil
    41  }
    42  
    43  type copyCmd struct{}
    44  
    45  func (cmd copyCmd) process(o *riofs.File, arg string) error {
    46  	log.Printf("copying %q...", arg)
    47  
    48  	fname, sel, err := splitArg(arg)
    49  	if err != nil {
    50  		return err
    51  	}
    52  	re := regexp.MustCompile(sel)
    53  
    54  	f, err := groot.Open(fname)
    55  	if err != nil {
    56  		return fmt.Errorf("could not open input ROOT file %q: %w", fname, err)
    57  	}
    58  	defer f.Close()
    59  
    60  	err = riofs.Walk(f, func(path string, obj root.Object, err error) error {
    61  		if err != nil {
    62  			return err
    63  		}
    64  		name := path[len(f.Name()):]
    65  		if !re.MatchString(name) {
    66  			return nil
    67  		}
    68  
    69  		var (
    70  			dst riofs.Directory
    71  			dir = stdpath.Dir(name)
    72  		)
    73  
    74  		odst, err := riofs.Dir(o).Get(dir)
    75  		if err != nil {
    76  			v, err := riofs.Dir(o).Mkdir(dir)
    77  			if err != nil {
    78  				return fmt.Errorf("could not create directory %q: %w", dir, err)
    79  			}
    80  			odst = v.(root.Object)
    81  		}
    82  		dst = odst.(riofs.Directory)
    83  
    84  		return cmd.copyObj(dst, stdpath.Base(name), obj)
    85  	})
    86  	if err != nil {
    87  		return fmt.Errorf("could not copy input ROOT file: %w", err)
    88  	}
    89  	return nil
    90  }
    91  
    92  func (cmd copyCmd) copyObj(odir riofs.Directory, k string, obj root.Object) error {
    93  	var err error
    94  	switch obj := obj.(type) {
    95  	case rtree.Tree:
    96  		err = cmd.copyTree(odir, k, obj)
    97  	case riofs.Directory:
    98  		_, err = odir.Mkdir(k)
    99  	default:
   100  		err = odir.Put(k, obj)
   101  	}
   102  
   103  	if err != nil {
   104  		return fmt.Errorf("could not save object %q to output file: %w", k, err)
   105  	}
   106  
   107  	return nil
   108  }
   109  
   110  func (cmd copyCmd) copyTree(dir riofs.Directory, name string, tree rtree.Tree) error {
   111  	dst, err := rtree.NewWriter(dir, name, rtree.WriteVarsFromTree(tree))
   112  	if err != nil {
   113  		return fmt.Errorf("could not create output copy tree: %w", err)
   114  	}
   115  	defer dst.Close()
   116  
   117  	src, err := rtree.NewReader(tree, nil)
   118  	if err != nil {
   119  		return fmt.Errorf("could not create tree reader: %w", err)
   120  	}
   121  	defer src.Close()
   122  
   123  	_, err = rtree.Copy(dst, src)
   124  	if err != nil {
   125  		return fmt.Errorf("could not copy tree %q: %w", name, err)
   126  	}
   127  
   128  	err = dst.Close()
   129  	if err != nil {
   130  		return fmt.Errorf("could not close copy tree %q: %w", name, err)
   131  	}
   132  
   133  	return nil
   134  }