import ROOT

ROOT.gInterpreter.Declare(
'''
#include "Math/Vector3D.h"
using namespace ROOT::Math;
''')

def test_warning():
    # Create two files to read data from
    # Warning appears only for cpp_vec and phys_vec branches
    for i in range(2):
        ROOT.RDataFrame(10).Define("simple_number", "rdfentry_")\
                           .Define("cpp_vec", "std::vector<double>{-1, 42, -1}")\
                           .Define("phys_vec", "XYZVector(-1, 42, -1)")\
                           .Snapshot("test_tree", "file{}.root".format(i))

    ch = ROOT.TChain("test_tree")
 
    # Using only one file shows No warning...
    # ch.Add("./file1.root")
    ch.Add("./file[01].root")

    df = ROOT.RDataFrame(ch).Filter('{cout<<"Event: "<<rdfentry_<< endl;} return true;')
    df = df.Define("new_var", "phys_vec.Y()*cpp_vec[1]")

    # Snapshoting w/o branch list shows no warning too
    # df.Snapshot("new_tree", "new_file.root")
    df.Snapshot("new_tree", "new_file.root", ["new_var"])


test_warning()