RootWriter.hh

Go to the documentation of this file.

#ifndef ROOTWRITER_h
#define ROOTWRITER_h

#include "BaseModule.hh"
#include <iostream>
#include <string>

//forward declarations
class TFile;
class TTree;
class RootWriter;

//utility classes
class BranchEnabler{
public:
  RootWriter* writer;
  BranchEnabler(RootWriter* w) : writer(w) {}
};

class BranchDisabler{
public:
  RootWriter* writer;
  BranchDisabler(RootWriter* w) : writer(w) {}
};

class RootWriter : public BaseModule{
public:
  static const std::string GetDefaultName(){ return "RootWriter"; }
  RootWriter();
  ~RootWriter();
  
  TTree* GetTree(){ return _tree; }
  int Initialize();
  int Finalize();
  int Process(EventPtr evt);
  
  const std::string GetFilename(){ return _filename; }
  void SetFilename(const std::string& name){ _filename=name; }
  static const std::string GetDefaultFilename(){ return "out.root"; }
  
  void EnableBranch(const char* classname, const char* branchname,
                    bool enable=true);
  void DisableBranch(const char* classname, const char* branchname);

private:
  void SaveConfig();
  std::string _filename;
  std::string _directory;
  std::string _mode;
  TFile* _outfile;
  TTree* _tree;
  bool default_saveall;
  BranchEnabler enabler;
  BranchDisabler disabler;
};

inline std::istream& operator>>(std::istream& in, BranchEnabler& e)
{
  std::string dummy;
  in>>dummy;
  size_t pos = dummy.find('.');
  if(pos == std::string::npos){
    Message(ERROR)<<"Branches to be enabled should follow the syntax CLASS.BRANCH.\n";
    return in;
  }
  e.writer->EnableBranch(dummy.substr(0,pos).c_str(), 
                          dummy.substr(pos+1).c_str());
  return in;
}

inline std::istream& operator>>(std::istream& in, BranchDisabler& d)
{
  std::string dummy;
  in>>dummy;
  size_t pos = dummy.find('.');
  if(pos == std::string::npos){
    Message(ERROR)<<"Branches to be disabled should follow the syntax CLASS.BRANCH.\n";
    return in;
  }
  d.writer->DisableBranch(dummy.substr(0,pos).c_str(), 
                          dummy.substr(pos+1).c_str());
  return in;
}

inline std::ostream& operator<<(std::ostream& out, const BranchEnabler& )
{ return out; }
inline std::ostream& operator<<(std::ostream&out, const BranchDisabler& )
{ return out; }

#endif