import ROOT
from array import array

def createTestFile( outFileName ):
    # Create output file and tree
    f = ROOT.TFile( outFileName, 'RECREATE')
    t = ROOT.TTree('MyTree','My test tree')
    
    # Create a branch and add to tree
    v =  array('f', [0.] )
    t.Branch( "v", v, "v" )                    
    
    w =  ROOT.std.vector( float )()
    setattr( t, "w", w )
    t.Branch( "w", w )            


    # Fill tree
    for i in range(10):
        v[0] = i
        w.resize(0)
        w.push_back(i)
        t.Fill()

    # Save and Close
    f.Write()
    f.Close()

# Create one file
createTestFile( "out-1.root" )

# Start local PROOF
ROOT.TProof.Open("workers=1")

# Create a TChain
chain = ROOT.TChain( "MyTree" );

# add the file with weight to the chain
for filename, weight in zip( ["out-1.root"], [10.]):
    f = ROOT.TFile.Open(filename,"UPDATE")
    tree = f.Get( "MyTree" )
    tree.SetWeight( weight )
    tree.AutoSave()
    f.Write()
    f.Close()
    chain.Add( filename )

# draw without PROOF
chain.Draw("v>>(100,0,10)")
ROOT.c1.Print("out-1.pdf")

# draw with PROOF
chain.SetProof()
chain.Draw("v>>(100,0,10)")
ROOT.c1.Print("out-2.pdf")

